feat: add new model

3q7y4e.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1646a218094821c8c0ca6df5c7f236bceb1aec6f4085d0a42f920bec6d53bb57
+size 352409020

app.py

@@ -3,6 +3,7 @@ import torch
 from PIL import Image
 import torchvision.transforms as transforms
 import numpy as np
+import torch.nn.functional as F
 from safetensors.torch import load_model, save_model
 from models import *
 import os
@@ -33,13 +34,24 @@ class WasteClassifier:
         img_tensor = self.transform(image).unsqueeze(0).to(self.device)
 
         with torch.no_grad():
-            outputs = self.model(img_tensor)
+            outputs, seg_mask = self.model(img_tensor)  # Handle both outputs
             probabilities = torch.nn.functional.softmax(outputs, dim=1)
 
         probs = probabilities[0].cpu().numpy()
         pred_class = self.class_names[np.argmax(probs)]
         confidence = np.max(probs)
 
+        # Process segmentation mask
+        seg_mask = (
+            seg_mask[0, 0].cpu().numpy().astype(np.float32)
+        )  # Get first image, first channel
+        # seg_mask = (seg_mask >= 0.2).astype(np.float32)  # Threshold at 0.2
+
+        # Resize mask back to original image size
+        seg_mask = Image.fromarray(seg_mask)
+        seg_mask = seg_mask.resize(original_size, Image.NEAREST)
+        seg_mask = np.array(seg_mask)
+
         results = {
             "predicted_class": pred_class,
             "confidence": confidence,
@@ -47,6 +59,7 @@ class WasteClassifier:
                 class_name: float(prob)
                 for class_name, prob in zip(self.class_names, probs)
             },
+            "segmentation_mask": seg_mask,
         }
 
         return results
@@ -56,6 +69,16 @@ def interface(classifier):
     def process_image(image):
         results = classifier.predict(image)
 
+        if isinstance(image, Image.Image):
+            image_np = np.array(image)
+        else:
+            image_np = image
+
+        mask = results["segmentation_mask"]
+
+        overlay = image_np.copy()
+        overlay[mask < 0.2] = overlay[mask < 0.2] * 0
+
         output_str = f"Predicted Class: {results['predicted_class']}\n"
         output_str += f"Confidence: {results['confidence']*100:.2f}%\n\n"
         output_str += "Class Probabilities:\n"
@@ -67,16 +90,23 @@ def interface(classifier):
         for class_name, prob in sorted_probs:
             output_str += f"{class_name}: {prob*100:.2f}%\n"
 
-        return output_str
+        mask_viz = (mask * 255).astype(np.uint8)
+
+        return [output_str, overlay, mask_viz]
 
     demo = gr.Interface(
         fn=process_image,
         inputs=[gr.Image(type="pil", label="Upload Image")],
-        outputs=[gr.Textbox(label="Classification Results")],
+        outputs=[
+            gr.Textbox(label="Classification Results"),
+            gr.Image(label="Segmented Object"),
+            gr.Image(label="Segmentation Mask"),
+        ],
         title="Waste Classification System",
         description="""
         Upload an image of waste to classify it into different categories.
-        The model will predict the type of waste and show confidence scores for each category.
+        The model will predict the type of waste, show confidence scores for each category,
+        and display the segmented object along with its mask.
         """,
         examples=(
             [["example1.jpg"], ["example2.jpg"], ["example3.jpg"]]
@@ -102,11 +132,12 @@ class_names = [
     "Textile Trash",
     "Vegetation",
 ]
-best_model = ResNet50(num_classes=len(class_names))
+
+best_model = ResNet101UNet(num_classes=len(class_names))
 best_model = best_model.to(device)
 load_model(
     best_model,
-    os.path.join(os.path.dirname(os.path.abspath(__file__)), "bjf8fp.safetensors"),
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "3q7y4e.safetensors"),
 )
 
 classifier = WasteClassifier(best_model, class_names, device)

models.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 
 class BasicBlock(nn.Module):
@@ -76,19 +77,20 @@ class Bottleneck(nn.Module):
 
 
 class ResNet(nn.Module):
-    def __init__(self, block, num_blocks, num_classes=1000, K=10, T=0.5):
+    def __init__(self, block, num_blocks, num_classes=1000):
         super(ResNet, self).__init__()
         self.in_planes = 64
-        self.K = K
-        self.T = T
 
         self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
         self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
         self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
         self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
         self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
         self.fc = nn.Linear(512 * block.expansion, num_classes)
 
     def _make_layer(self, block, planes, num_blocks, stride):
@@ -99,43 +101,201 @@ class ResNet(nn.Module):
             self.in_planes = planes * block.expansion
         return nn.Sequential(*layers)
 
-    def t_max_avg_pooling(self, x):
-        B, C, H, W = x.shape
-        x_flat = x.view(B, C, -1)
-        top_k_values, _ = torch.topk(x_flat, self.K, dim=2)
-        max_values = top_k_values.max(dim=2)[0]
-        avg_values = top_k_values.mean(dim=2)
-        output = torch.where(max_values >= self.T, max_values, avg_values)
-        return output
-
     def forward(self, x):
         out = torch.relu(self.bn1(self.conv1(x)))
         out = self.maxpool(out)
+
         out = self.layer1(out)
         out = self.layer2(out)
         out = self.layer3(out)
         out = self.layer4(out)
-        out = self.t_max_avg_pooling(out)
-        out = out.view(out.size(0), -1)
+
+        out = self.avgpool(out)
+        out = torch.flatten(out, 1)
         out = self.fc(out)
         return out
 
 
-def ResNet18(num_classes=1000, K=10, T=0.5):
-    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes, K, T)
+def ResNet18(num_classes=1000):
+    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes)
+
+
+def ResNet34(num_classes=1000):
+    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes)
+
+
+def ResNet50(num_classes=1000):
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes)
+
+
+def ResNet101(num_classes=1000):
+    return ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
+
+
+def ResNet152(num_classes=1000):
+    return ResNet(Bottleneck, [3, 8, 36, 3], num_classes)
+
+
+class ClassifierHead(nn.Module):
+    def __init__(self, in_features, num_classes):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.max_pool = nn.AdaptiveMaxPool2d((1, 1))
+
+        self.classifier = nn.Sequential(
+            nn.Linear(in_features * 2, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes),
+        )
+
+    def forward(self, x):
+        avg_pooled = self.avg_pool(x).flatten(1)
+        max_pooled = self.max_pool(x).flatten(1)
+        features = torch.cat([avg_pooled, max_pooled], dim=1)
+        return self.classifier(features)
+
+
+class ResNetUNet(ResNet):
+    def __init__(self, block, num_blocks, num_classes=1000):
+        super().__init__(block, num_blocks, num_classes)
+
+        # Calculate encoder channel sizes
+        self.enc_channels = [
+            64,
+            64 * block.expansion,
+            128 * block.expansion,
+            256 * block.expansion,
+            512 * block.expansion,
+        ]
+
+        # Replace t_max_avg_pooling with standard avgpool
+        in_features = 512 * block.expansion
+        self.classifier_head = ClassifierHead(in_features, num_classes)
+
+        # Decoder layers remain the same
+        self.decoder5 = nn.Sequential(
+            nn.Conv2d(2048 + 1024, 1024, 3, padding=1),
+            nn.BatchNorm2d(1024),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024, 512, 3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+        )
+
+        self.decoder4 = nn.Sequential(
+            nn.Conv2d(512 + 512, 512, 3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 256, 3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+        )
+
+        self.decoder3 = nn.Sequential(
+            nn.Conv2d(256 + 256, 256, 3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 128, 3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+        )
+
+        self.decoder2 = nn.Sequential(
+            nn.Conv2d(128 + 64, 128, 3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 64, 3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+        )
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(64, 32, 3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, 1, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        input_size = x.shape[-2:]
+
+        # Encoder path
+        x = torch.relu(self.bn1(self.conv1(x)))
+        e1 = self.maxpool(x)
+
+        e2 = self.layer1(e1)
+        e3 = self.layer2(e2)
+        e4 = self.layer3(e3)
+        e5 = self.layer4(e4)
+
+        # Get segmentation first
+        e4_resized = F.interpolate(
+            e4, size=e5.shape[-2:], mode="bilinear", align_corners=True
+        )
+        d5 = self.decoder5(torch.cat([e5, e4_resized], dim=1))
+
+        e3_resized = F.interpolate(
+            e3, size=d5.shape[-2:], mode="bilinear", align_corners=True
+        )
+        d4 = self.decoder4(torch.cat([d5, e3_resized], dim=1))
+
+        e2_resized = F.interpolate(
+            e2, size=d4.shape[-2:], mode="bilinear", align_corners=True
+        )
+        d3 = self.decoder3(torch.cat([d4, e2_resized], dim=1))
+
+        e1_resized = F.interpolate(
+            e1, size=d3.shape[-2:], mode="bilinear", align_corners=True
+        )
+        d2 = self.decoder2(torch.cat([d3, e1_resized], dim=1))
+
+        seg_out = self.final_conv(d2)
+        seg_out = F.interpolate(
+            seg_out, size=input_size, mode="bilinear", align_corners=True
+        )
+
+        # Use segmentation to mask features before classification
+        # Upsample segmentation mask to match feature size
+        attention_mask = F.interpolate(
+            seg_out, size=e5.shape[2:], mode="bilinear", align_corners=True
+        )
+
+        # Apply attention mask to features
+        attended_features = e5 * (0.25 + attention_mask)
+
+        # Use new classifier head
+        cls_out = self.classifier_head(attended_features)
+
+        return cls_out, seg_out
+
+
+# Helper functions without K and T parameters
+def ResNet18UNet(num_classes=1000):
+    return ResNetUNet(BasicBlock, [2, 2, 2, 2], num_classes)
 
 
-def ResNet34(num_classes=1000, K=10, T=0.5):
-    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes, K, T)
+def ResNet34UNet(num_classes=1000):
+    return ResNetUNet(BasicBlock, [3, 4, 6, 3], num_classes)
 
 
-def ResNet50(num_classes=1000, K=10, T=0.5):
-    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes, K, T)
+def ResNet50UNet(num_classes=1000):
+    return ResNetUNet(Bottleneck, [3, 4, 6, 3], num_classes)
 
 
-def ResNet101(num_classes=1000, K=10, T=0.5):
-    return ResNet(Bottleneck, [3, 4, 23, 3], num_classes, K, T)
+def ResNet101UNet(num_classes=1000):
+    return ResNetUNet(Bottleneck, [3, 4, 23, 3], num_classes)
 
 
-def ResNet152(num_classes=1000, K=10, T=0.5):
-    return ResNet(Bottleneck, [3, 8, 36, 3], num_classes, K, T)
+def ResNet152UNet(num_classes=1000):
+    return ResNetUNet(Bottleneck, [3, 8, 36, 3], num_classes)