Upd FastAPI backend img processing garbage det and cls + frontend JS side rendering

app.py

@@ -40,6 +40,7 @@ os.environ["HF_HOME"]           = CACHE
 
 # ── Load models once  ───────────────────────────────────────────────────
 print("🔄  Loading models …")
+# Garbage detection
 model_self  = YOLO(f"{MODEL_DIR}/garbage_detector.pt")                 # YOLOv11(l)
 model_yolo5 = yolov5.load(f"{MODEL_DIR}/yolov5-detect-trash-classification.pt")
 processor_detr = DetrImageProcessor.from_pretrained(f"{MODEL_DIR}/detr")
@@ -48,7 +49,10 @@ feat_extractor = SegformerFeatureExtractor.from_pretrained(
                     "nvidia/segformer-b4-finetuned-ade-512-512")
 segformer      = SegformerForSemanticSegmentation.from_pretrained(
                     "nvidia/segformer-b4-finetuned-ade-512-512")
+# Animal detection
 model_animal = YOLO(f"{MODEL_DIR}/yolov8n.pt") # Load COCO pre-trained YOLOv8 for animal detection
+# Garbage classification
+model_garbage_cls = YOLO(f"{MODEL_DIR}/garbage_cls_yolov8s.pt") 
 print("✅  Models ready\n")
 
 # ── ADE-20K palette + custom mapping (verbatim) ─────────────────────────
@@ -413,6 +417,55 @@ async def detect_animals(file: UploadFile = File(...)):
     return FileResponse(result_path, media_type="image/jpeg")
 
 
+# Garbage classification endpoint
+@app.post("/classification/")
+async def classify_garbage(file: UploadFile = File(...)):
+    img_id = _uid()
+    img_path = f"{UPLOAD_DIR}/{img_id}_{file.filename}"
+    out_path = f"{OUTPUT_DIR}/{img_id}_classified.jpg"
+    # Load file
+    with open(img_path, "wb") as f:
+        shutil.copyfileobj(file.file, f)
+    # Read image
+    print(f"[Classification] Received image: {img_path}")
+    image = cv2.imread(img_path)
+    rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    pil = Image.fromarray(rgb)
+    # DETR for garbage detection boxes
+    detections = []
+    inp = processor_detr(images=pil, return_tensors="pt")
+    with torch.no_grad():
+        out = model_detr(**inp)
+    results = processor_detr.post_process_object_detection(
+        outputs=out,
+        target_sizes=torch.tensor([pil.size[::-1]]),
+        threshold=0.5
+    )[0]
+    # Bbox return
+    boxes = results["boxes"]
+    print(f"[Classification] {len(boxes)} garbage objects detected by DETR.")
+    # Mapping in between
+    for i, box in enumerate(boxes):
+        x1, y1, x2, y2 = map(int, box.tolist())
+        crop = image[y1:y2, x1:x2]
+        if crop.shape[0] < 10 or crop.shape[1] < 10:
+            continue  # skip tiny crops
+        # Convert crop to RGB and classify
+        pil_crop = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
+        pred = model_garbage_cls(pil_crop, verbose=False)[0]
+        class_id = int(pred.probs.top1)
+        class_name = model_garbage_cls.names[class_id]
+        conf = pred.probs.top1conf
+        # Labelling on output image
+        label = f"{class_name} ({conf:.2f})"
+        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 165, 255), 2)
+        cv2.putText(image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 165, 255), 2)
+    # Write image on render
+    cv2.imwrite(out_path, image)
+    print(f"[Classification] Output saved: {out_path}")
+    return FileResponse(out_path, media_type="image/jpeg")
+
+
 # ── Core pipeline (runs in background thread) ───────────────────────────
 def _pipeline(uid,img_path):
     print(f"▶️ [{uid}] processing")

statics/index.html

@@ -24,7 +24,7 @@
     <div id="animal-result"></div>
     <div id="upload-container3">
         <input type="file" id="upload3" accept="image/*">
-        <button  id="checkTrashBtn" onclick="uploadAnimal()">Check Animal</button>
+        <button  id="checkTrashBtn" onclick="uploadTrash()">Classify Garbage</button>
     </div>
     <div id="trash-result"></div>
     <script src="/statics/script.js"></script>

statics/script.js

@@ -62,5 +62,29 @@ async function uploadAnimal() {
     const imgURL = URL.createObjectURL(blob);
     document.getElementById("animal-result").innerHTML =
       `<p><b>Animal Detection Result:</b></p><img src="${imgURL}" width="640"/>`;
+}
+
+async function uploadTrash() {
+    const fileInput = document.getElementById('upload3');
+    if (!fileInput.files.length) return alert("Upload an image first");
+    // Upload and read image file
+    const formData = new FormData();
+    formData.append("file", fileInput.files[0]);
+    // Handshake with FastAPI
+    const res = await fetch("/classification/", {
+      method: "POST",
+      body: formData
+    });
+    // Error
+    if (!res.ok) {
+      alert("Failed to process garbage classification.");
+      return;
+    }
+    // Create image
+    const blob = await res.blob();
+    const imgURL = URL.createObjectURL(blob);
+    document.getElementById("trash-result").innerHTML =
+      `<p><b>Garbage Classification Result:</b></p><img src="${imgURL}" width="640"/>`;
   }
+