Spaces:
Runtime error
Runtime error
| import base64 | |
| import io | |
| import math | |
| import os | |
| import time | |
| try: | |
| import spaces | |
| except ImportError: | |
| class _SpacesFallback: | |
| def GPU(*_args, **_kwargs): | |
| def decorator(func): | |
| return func | |
| return decorator | |
| spaces = _SpacesFallback() | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from fastapi import FastAPI, File, HTTPException, UploadFile | |
| from fastapi.responses import JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from PIL import Image | |
| from transformers import ViTForImageClassification, ViTImageProcessor | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def load_model_safely(): | |
| """Load model with fallback options and proper error handling. | |
| Returns (processor, model, model_id) — model_id names the source that | |
| actually loaded, so the analysis panel reports the real weights in use. | |
| """ | |
| try: | |
| if os.path.exists("./waste-classification"): | |
| print("Loading model from local clone...") | |
| processor = ViTImageProcessor.from_pretrained("./waste-classification") | |
| model = ViTForImageClassification.from_pretrained( | |
| "./waste-classification", attn_implementation="eager" | |
| ) | |
| print("Successfully loaded model from local clone") | |
| return processor, model, "local: ./waste-classification" | |
| except Exception as e: | |
| print(f"Failed to load from local clone: {e}") | |
| try: | |
| print("Loading model from HuggingFace...") | |
| processor = ViTImageProcessor.from_pretrained( | |
| "watersplash/waste-classification", cache_dir="./cache" | |
| ) | |
| # eager attention so /api/explain can read attention weights (sdpa cannot) | |
| model = ViTForImageClassification.from_pretrained( | |
| "watersplash/waste-classification", cache_dir="./cache", attn_implementation="eager" | |
| ) | |
| print("Successfully loaded model from HuggingFace") | |
| return processor, model, "watersplash/waste-classification" | |
| except Exception as e: | |
| print(f"Failed to load from HuggingFace: {e}") | |
| try: | |
| print("Loading base ViT model as fallback...") | |
| processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k") | |
| model = ViTForImageClassification.from_pretrained( | |
| "google/vit-base-patch16-224-in21k", | |
| attn_implementation="eager", | |
| num_labels=12, | |
| id2label={ | |
| "0": "battery", "1": "biological", "2": "brown-glass", "3": "cardboard", | |
| "4": "clothes", "5": "green-glass", "6": "metal", "7": "paper", | |
| "8": "plastic", "9": "shoes", "10": "trash", "11": "white-glass", | |
| }, | |
| label2id={ | |
| "battery": "0", "biological": "1", "brown-glass": "2", "cardboard": "3", | |
| "clothes": "4", "green-glass": "5", "metal": "6", "paper": "7", | |
| "plastic": "8", "shoes": "9", "trash": "10", "white-glass": "11", | |
| }, | |
| ) | |
| print("Loaded base ViT model as fallback (untrained)") | |
| return processor, model, "google/vit-base-patch16-224-in21k (untrained fallback)" | |
| except Exception as e: | |
| print(f"Failed to load fallback model: {e}") | |
| return None, None, "none" | |
| print("Initializing model...") | |
| processor, model, MODEL_ID = load_model_safely() | |
| if model is not None: | |
| model.to(DEVICE) | |
| model.eval() | |
| CLASS_NAMES = [ | |
| "Battery", "Biological", "Brown-glass", "Cardboard", "Clothes", | |
| "Green-glass", "Metal", "Paper", "Plastic", "Shoes", "Trash", "White-glass", | |
| ] | |
| # Waste segregation mapping: which bin each class belongs to | |
| BIN_MAP = { | |
| "Biological": "green", | |
| "Cardboard": "blue", "Paper": "blue", "Plastic": "blue", "Metal": "blue", | |
| "Brown-glass": "blue", "Green-glass": "blue", "White-glass": "blue", | |
| "Clothes": "blue", "Shoes": "blue", | |
| "Battery": "red", "Trash": "red", | |
| } | |
| BIN_INFO = { | |
| "green": {"name": "Green Bin", "desc": "Wet / Biodegradable Waste"}, | |
| "blue": {"name": "Blue Bin", "desc": "Dry / Recyclable Waste"}, | |
| "red": {"name": "Red Bin", "desc": "Hazardous / Reject Waste"}, | |
| } | |
| def _classify_waste_impl(image: Image.Image): | |
| """Classify a PIL image, return (label, confidence, bin_color, top, analysis)""" | |
| if processor is None or model is None: | |
| raise RuntimeError("Model failed to load") | |
| if image.mode != "RGB": | |
| image = image.convert("RGB") | |
| original_size = list(image.size) | |
| started = time.perf_counter() | |
| inputs = processor(images=image, return_tensors="pt") | |
| inputs = {name: tensor.to(DEVICE) for name, tensor in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| predictions = torch.nn.functional.softmax(outputs.logits, dim=-1) | |
| latency_ms = (time.perf_counter() - started) * 1000 | |
| confidence_scores = predictions[0].tolist() | |
| ranked = sorted(range(len(confidence_scores)), key=lambda i: confidence_scores[i], reverse=True) | |
| best_idx = ranked[0] | |
| label = CLASS_NAMES[best_idx] | |
| confidence = confidence_scores[best_idx] | |
| bin_color = BIN_MAP.get(label, "red") | |
| top = [ | |
| {"label": CLASS_NAMES[i], "confidence": round(confidence_scores[i], 4)} | |
| for i in ranked[:3] | |
| ] | |
| # Shannon entropy in bits; normalised against log2(n) so 0 = certain, 1 = uniform. | |
| entropy = -sum(p * math.log2(p) for p in confidence_scores if p > 0) | |
| max_entropy = math.log2(len(confidence_scores)) | |
| margin = confidence_scores[ranked[0]] - confidence_scores[ranked[1]] | |
| analysis = { | |
| "distribution": [ | |
| {"label": CLASS_NAMES[i], "confidence": round(confidence_scores[i], 6)} | |
| for i in ranked | |
| ], | |
| "entropy_bits": round(entropy, 4), | |
| "entropy_normalized": round(entropy / max_entropy, 4) if max_entropy else 0.0, | |
| "margin": round(margin, 4), | |
| "logit_top": round(outputs.logits[0][best_idx].item(), 4), | |
| "latency_ms": round(latency_ms, 1), | |
| "device": str(next(model.parameters()).device), | |
| "model": { | |
| "id": MODEL_ID, | |
| "architecture": model.config.architectures[0] if model.config.architectures else "ViTForImageClassification", | |
| "parameters": sum(p.numel() for p in model.parameters()), | |
| "input_size": [model.config.image_size, model.config.image_size], | |
| "patch_size": model.config.patch_size, | |
| "hidden_size": model.config.hidden_size, | |
| "num_layers": model.config.num_hidden_layers, | |
| "num_heads": model.config.num_attention_heads, | |
| "num_classes": len(CLASS_NAMES), | |
| "original_image_size": original_size, | |
| "normalize_mean": list(processor.image_mean), | |
| "normalize_std": list(processor.image_std), | |
| }, | |
| } | |
| return label, confidence, bin_color, top, analysis | |
| def classify_waste(image: Image.Image): | |
| return _classify_waste_impl(image) | |
| # --- FastAPI app: serves the vanilla JS frontend + classification API --- | |
| app = FastAPI(title="Waste Classification") | |
| async def api_classify(file: UploadFile = File(...)): | |
| if processor is None or model is None: | |
| raise HTTPException(status_code=503, detail="Model failed to load. Please try again later.") | |
| try: | |
| image_bytes = await file.read() | |
| image = Image.open(io.BytesIO(image_bytes)) | |
| except Exception: | |
| raise HTTPException(status_code=400, detail="Invalid image file.") | |
| try: | |
| label, confidence, bin_color, top, analysis = classify_waste(image) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Classification failed: {e}") | |
| info = BIN_INFO[bin_color] | |
| return JSONResponse({ | |
| "label": label, | |
| "confidence": round(confidence, 4), | |
| "bin_color": bin_color, | |
| "bin_name": info["name"], | |
| "bin_desc": info["desc"], | |
| "top": top, | |
| "analysis": analysis, | |
| }) | |
| def _colorize(heat): | |
| """Map a 0..1 attention map to a blue→cyan→green→yellow→red ramp.""" | |
| stops = np.array([ | |
| [59, 76, 192], [77, 165, 220], [92, 200, 160], | |
| [237, 205, 74], [214, 86, 51], | |
| ], dtype=np.float32) | |
| pos = np.clip(heat, 0.0, 1.0) * (len(stops) - 1) | |
| lo = np.floor(pos).astype(int) | |
| hi = np.clip(lo + 1, 0, len(stops) - 1) | |
| t = (pos - lo)[..., None] | |
| return (stops[lo] * (1 - t) + stops[hi] * t).astype(np.uint8) | |
| def attention_rollout(image: Image.Image): | |
| """Overlay ViT attention rollout on the image; returns a PNG data URL. | |
| Rollout follows Abnar & Zuidema (2020): average heads per layer, add the | |
| residual connection as an identity, renormalise, then multiply across | |
| layers to trace CLS attention back to the input patches. | |
| """ | |
| if processor is None or model is None: | |
| raise RuntimeError("Model failed to load") | |
| if image.mode != "RGB": | |
| image = image.convert("RGB") | |
| started = time.perf_counter() | |
| inputs = processor(images=image, return_tensors="pt") | |
| inputs = {name: tensor.to(DEVICE) for name, tensor in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model(**inputs, output_attentions=True) | |
| rollout = None | |
| for layer in outputs.attentions: | |
| attn = layer[0].mean(dim=0) | |
| attn = attn + torch.eye(attn.size(0), device=attn.device) | |
| attn = attn / attn.sum(dim=-1, keepdim=True) | |
| rollout = attn if rollout is None else attn @ rollout | |
| grid = model.config.image_size // model.config.patch_size | |
| heat = rollout[0, 1:].reshape(grid, grid).detach().cpu().numpy() | |
| # Rollout is diffuse by nature; stretch the informative band so the | |
| # attended region reads clearly instead of a flat wash. | |
| lo, hi = np.percentile(heat, 55), np.percentile(heat, 99.5) | |
| heat = np.clip((heat - lo) / (hi - lo + 1e-8), 0.0, 1.0) ** 1.3 | |
| # Render at most 512px on the long side to keep the payload small. | |
| preview = image.copy() | |
| preview.thumbnail((512, 512), Image.LANCZOS) | |
| heat_img = Image.fromarray((heat * 255).astype(np.uint8)).resize(preview.size, Image.BICUBIC) | |
| heat_arr = np.asarray(heat_img, dtype=np.float32) / 255.0 | |
| base = np.asarray(preview, dtype=np.float32) | |
| alpha = (0.72 * heat_arr)[..., None] | |
| blended = base * (1 - alpha) + _colorize(heat_arr).astype(np.float32) * alpha | |
| buffer = io.BytesIO() | |
| Image.fromarray(blended.astype(np.uint8)).save(buffer, format="PNG", optimize=True) | |
| encoded = base64.b64encode(buffer.getvalue()).decode("ascii") | |
| return { | |
| "heatmap": f"data:image/png;base64,{encoded}", | |
| "grid": [grid, grid], | |
| "patches": grid * grid, | |
| "layers": len(outputs.attentions), | |
| "latency_ms": round((time.perf_counter() - started) * 1000, 1), | |
| } | |
| async def api_explain(file: UploadFile = File(...)): | |
| if processor is None or model is None: | |
| raise HTTPException(status_code=503, detail="Model failed to load. Please try again later.") | |
| try: | |
| image_bytes = await file.read() | |
| image = Image.open(io.BytesIO(image_bytes)) | |
| except Exception: | |
| raise HTTPException(status_code=400, detail="Invalid image file.") | |
| try: | |
| return JSONResponse(attention_rollout(image)) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Explanation failed: {e}") | |
| async def health(): | |
| return {"status": "ok" if (model is not None and processor is not None) else "model_not_loaded"} | |
| # --- Minimal Gradio Blocks, mounted at /gradio purely to satisfy the Gradio SDK --- | |
| def _gradio_classify(image): | |
| if image is None: | |
| return {"Please upload an image": 1.0} | |
| try: | |
| label, confidence, bin_color, _top, _analysis = classify_waste(image) | |
| info = BIN_INFO[bin_color] | |
| return {f"{label} → {info['name']}": confidence} | |
| except Exception as e: | |
| return {f"Error: {e}": 1.0} | |
| with gr.Blocks(title="Waste Classification (debug UI)") as demo: | |
| gr.Markdown("### Debug UI — the real app lives at the Space root URL") | |
| gr.Interface( | |
| fn=_gradio_classify, | |
| inputs=gr.Image(type="pil", label="Upload Waste Image"), | |
| outputs=gr.Textbox(label="Result"), | |
| ) | |
| # Serve the vanilla JS frontend at the root when the FastAPI app is used outside ZeroGPU. | |
| app.mount("/", StaticFiles(directory="static", html=True), name="static") | |
| def _zerogpu_classify(image): | |
| if image is None: | |
| return "Please upload an image." | |
| label, confidence, bin_color, _top, _analysis = _classify_waste_impl(image) | |
| info = BIN_INFO[bin_color] | |
| return ( | |
| f"Prediction: {label}\n" | |
| f"Confidence: {confidence:.1%}\n" | |
| f"Bin: {info['name']}\n" | |
| f"Description: {info['desc']}" | |
| ) | |
| with gr.Blocks(title="Waste Classification") as zerogpu_demo: | |
| gr.Markdown("### Waste Classification") | |
| zerogpu_image = gr.Image(type="pil", label="Upload Waste Image") | |
| zerogpu_button = gr.Button("Classify") | |
| zerogpu_result = gr.Textbox(label="Result") | |
| zerogpu_button.click( | |
| fn=_zerogpu_classify, | |
| inputs=zerogpu_image, | |
| outputs=zerogpu_result, | |
| api_name="classify", | |
| ) | |
| if __name__ == "__main__": | |
| zerogpu_demo.launch(server_name="0.0.0.0", server_port=7860) | |