feat: add project files

.gitignore

@@ -0,0 +1 @@
+__pycache__/

app.py

@@ -0,0 +1,112 @@
+import gradio as gr
+import torch
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+from safetensors.torch import load_model, save_model
+from models import *
+import os
+
+
+class WasteClassifier:
+    def __init__(self, model, class_names, device):
+        self.model = model
+        self.class_names = class_names
+        self.device = device
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize((384, 384)),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    def predict(self, image):
+        self.model.eval()
+
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        original_size = image.size
+        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model(img_tensor)
+            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)
+
+        results = {
+            "predicted_class": pred_class,
+            "confidence": confidence,
+            "class_probabilities": {
+                class_name: float(prob)
+                for class_name, prob in zip(self.class_names, probs)
+            },
+        }
+
+        return results
+
+
+def interface(classifier):
+    def process_image(image):
+        results = classifier.predict(image)
+
+        output_str = f"Predicted Class: {results['predicted_class']}\n"
+        output_str += f"Confidence: {results['confidence']*100:.2f}%\n\n"
+        output_str += "Class Probabilities:\n"
+
+        sorted_probs = sorted(
+            results["class_probabilities"].items(), key=lambda x: x[1], reverse=True
+        )
+
+        for class_name, prob in sorted_probs:
+            output_str += f"{class_name}: {prob*100:.2f}%\n"
+
+        return output_str
+
+    demo = gr.Interface(
+        fn=process_image,
+        inputs=[gr.Image(type="pil", label="Upload Image")],
+        outputs=[gr.Textbox(label="Classification Results")],
+        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.
+        """,
+        examples=(
+            [["example1.jpg"], ["example2.jpg"], ["example3.jpg"]]
+            if os.path.exists("example1.jpg")
+            else None
+        ),
+        analytics_enabled=False,
+        theme="default",
+    )
+
+    return demo
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+class_names = [
+    "Cardboard",
+    "Food Organics",
+    "Glass",
+    "Metal",
+    "Miscellaneous Trash",
+    "Paper",
+    "Plastic",
+    "Textile Trash",
+    "Vegetation",
+]
+best_model = ResNet50(num_classes=len(class_names))
+best_model = best_model.to(device)
+load_model(best_model, os.path.join(__file__, "..", "bjf8fp.safetensors"))
+
+classifier = WasteClassifier(best_model, class_names, device)
+
+demo = interface(classifier)
+demo.launch(share=True)

bjf8fp.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94e364935fe6427d14f6c9aaaa179ab707d33b27d1a3cf4705e6033595e562ea
+size 94347292

models.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, stride=1, padding=1, bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = torch.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = torch.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, self.expansion * planes, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = torch.relu(self.bn1(self.conv1(x)))
+        out = torch.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = torch.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=1000, K=10, T=0.5):
+        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.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            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.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 ResNet34(num_classes=1000, K=10, T=0.5):
+    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes, K, T)
+
+
+def ResNet50(num_classes=1000, K=10, T=0.5):
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes, K, T)
+
+
+def ResNet101(num_classes=1000, K=10, T=0.5):
+    return ResNet(Bottleneck, [3, 4, 23, 3], num_classes, K, T)
+
+
+def ResNet152(num_classes=1000, K=10, T=0.5):
+    return ResNet(Bottleneck, [3, 8, 36, 3], num_classes, K, T)

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+pillow
+gradio
+numpy