Update classifier.py

classifier.py

@@ -1,99 +1,99 @@
-import torch
-from torchvision import models, transforms
-from PIL import Image
-
-from waste_logic import map_to_waste, get_explanation
-
-# -------------------------
-# Device setup
-# -------------------------
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# -------------------------
-# Lazy-loaded model (important for backend)
-# -------------------------
-_model = None
-
-def get_model():
-    global _model
-    if _model is None:
-        model = models.resnet50(weights=None)
-        state_dict = torch.load("resnet50.pth", map_location="cpu")
-        model.load_state_dict(state_dict)
-        model.to(device)
-        model.eval()
-        _model = model
-    return _model
-
-# -------------------------
-# Image preprocessing
-# -------------------------
-preprocess = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    )
-])
-
-# -------------------------
-# Load ImageNet labels
-# -------------------------
-with open("imagenet_classes.txt", "r") as f:
-    labels = [line.strip() for line in f.readlines()]
-
-# -------------------------
-# Image classification (Top-K)
-# -------------------------
-def classify_image(image_path, top_k=3):
-    model = get_model()
-
-    try:
-        image = Image.open(image_path).convert("RGB")
-    except Exception as e:
-        raise ValueError(f"Invalid image file: {e}")
-
-    tensor = preprocess(image).unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        outputs = model(tensor)
-        probs = torch.softmax(outputs, dim=1)
-
-    top_probs, top_idxs = torch.topk(probs, top_k)
-
-    results = []
-    for prob, idx in zip(top_probs[0], top_idxs[0]):
-        results.append((labels[idx.item()], prob.item()))
-
-    return results
-
-# -------------------------
-# Public API function
-# -------------------------
-def analyze_image(image_path):
-    predictions = classify_image(image_path)
-
-    chosen_label = predictions[0][0]
-    chosen_conf = predictions[0][1]
-    waste_type = "Unknown"
-
-    for label, conf in predictions:
-        wt = map_to_waste(label)
-        if wt != "Unknown":
-            chosen_label = label
-            chosen_conf = conf
-            waste_type = wt
-            break
-
-    explanation = get_explanation(waste_type)
-
-    return {
-        "detected_object": chosen_label,
-        "confidence": round(chosen_conf, 3),
-        "waste_category": waste_type,
-        "reason": explanation["reason"],
-        "disposal": explanation["disposal"],
-        "environmental_impact": explanation["impact"]
-    }
+import torch
+from torchvision import models, transforms
+from PIL import Image
+
+from waste_logic import map_to_waste, get_explanation
+
+# -------------------------
+# Device setup
+# -------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# -------------------------
+# Lazy-loaded model (important for backend)
+# -------------------------
+_model = None
+
+def get_model():
+    global _model
+    if _model is None:
+        model = models.resnet50(weights=None)
+        state_dict = torch.load("resnet50.pth", map_location=device)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        _model = model
+    return _model
+
+# -------------------------
+# Image preprocessing
+# -------------------------
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+# -------------------------
+# Load ImageNet labels
+# -------------------------
+with open("imagenet_classes.txt", "r") as f:
+    labels = [line.strip() for line in f.readlines()]
+
+# -------------------------
+# Image classification (Top-K)
+# -------------------------
+def classify_image(image_path, top_k=3):
+    model = get_model()
+
+    try:
+        image = Image.open(image_path).convert("RGB")
+    except Exception as e:
+        raise ValueError(f"Invalid image file: {e}")
+
+    tensor = preprocess(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(tensor)
+        probs = torch.softmax(outputs, dim=1)
+
+    top_probs, top_idxs = torch.topk(probs, top_k)
+
+    results = []
+    for prob, idx in zip(top_probs[0], top_idxs[0]):
+        results.append((labels[idx.item()], prob.item()))
+
+    return results
+
+# -------------------------
+# Public API function
+# -------------------------
+def analyze_image(image_path):
+    predictions = classify_image(image_path)
+
+    chosen_label = predictions[0][0]
+    chosen_conf = predictions[0][1]
+    waste_type = "Unknown"
+
+    for label, conf in predictions:
+        wt = map_to_waste(label)
+        if wt != "Unknown":
+            chosen_label = label
+            chosen_conf = conf
+            waste_type = wt
+            break
+
+    explanation = get_explanation(waste_type)
+
+    return {
+        "detected_object": chosen_label,
+        "confidence": round(chosen_conf, 3),
+        "waste_category": waste_type,
+        "reason": explanation["reason"],
+        "disposal": explanation["disposal"],
+        "environmental_impact": explanation["impact"]
+    }