import os import io import json import base64 import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image import cv2 import numpy as np import gradio as gr from transformers import AutoModelForImageClassification, AutoImageProcessor # --- Configuration --- # Use an environment variable for the token to stay secure HF_TOKEN = os.environ.get("HF_TOKEN") USE_HF_TRANSFORMERS = True HF_MODEL_ID = "kritaphatson/thai_snake_image_classifier" # Use repo ID for Space deployment # Determine model path (prioritize local if exists, else Hub) LOCAL_MODEL_PATH = r'C:\Users\ADMIN\PycharmProjects\snake_research\transformer_experiment_results_swin_exact\Swin-Base_augnone_fineT_thres500' model_path = LOCAL_MODEL_PATH if os.path.exists(LOCAL_MODEL_PATH) else HF_MODEL_ID # Grad-CAM Settings STAGE_INDEX = 2 AGG_METHOD = "layercam" ALPHA = 0.40 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def load_inference_engine(): try: print(f"Attempting to load model from: {model_path}") # Use token for both local and Hub to be safe model = AutoModelForImageClassification.from_pretrained(model_path, token=HF_TOKEN) processor = AutoImageProcessor.from_pretrained(model_path, token=HF_TOKEN) model.to(device) model.eval() print("Model loaded successfully.") except Exception as e: print(f"Error loading model: {e}") return None, None, {} # Try to load id2label idx_to_class = {} if hasattr(model.config, 'id2label') and model.config.id2label: idx_to_class = model.config.id2label print("Using model.config.id2label") # Check for local file in Space dir if not idx_to_class: local_id2label = "id2label.json" if os.path.exists(local_id2label): with open(local_id2label, 'r') as f: idx_to_class = json.load(f) print(f"Loaded id2label from local {local_id2label}") if not idx_to_class: print("Warning: No id2label found, using defaults.") idx_to_class = {str(i): f"Class {i}" for i in range(model.config.num_labels)} return model, processor, idx_to_class model, processor, idx_to_class = load_inference_engine() str_idx_to_class = {str(k): v for k, v in idx_to_class.items()} if idx_to_class else {} # --- Grad-CAM Logic --- # (Same logic as before, but with safety) def get_gradcam(model, pixel_values, target_layer): if model is None: return None holder = {"acts": None, "grads": None} def fwd_hook(_m, _inp, out): hs = out[0] if isinstance(out, (tuple, list)) else out hs.requires_grad_(True) hs.retain_grad() holder["acts"] = hs hs.register_hook(lambda g: holder.__setitem__("grads", g)) hook = target_layer.register_forward_hook(fwd_hook) model.zero_grad() outputs = model(pixel_values=pixel_values) logits = outputs.logits idx = torch.argmax(logits, dim=-1) score = logits[0, idx] score.backward() hook.remove() acts = holder["acts"][0].detach().cpu() grads = holder["grads"][0].detach().cpu() if AGG_METHOD == "layercam": cam = torch.relu(grads) * torch.relu(acts) cam = cam.sum(dim=-1) else: weights = grads.mean(dim=0) cam = torch.relu(acts @ weights) cam = cam - cam.min() if cam.max() > 0: cam = cam / cam.max() s = int(np.sqrt(cam.shape[0])) return cam.reshape(s, s).numpy() def apply_overlay(orig_img, heatmap, alpha=ALPHA): heatmap = cv2.resize(heatmap, (orig_img.size[0], orig_img.size[1])) heatmap = np.uint8(255 * heatmap) heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) overlayed = cv2.addWeighted(heatmap, alpha, np.array(orig_img), 1 - alpha, 0) return Image.fromarray(overlayed) def predict(image): if image is None: return None, {"Error": "No image uploaded"} if model is None or processor is None: return None, {"Error": "Model failed to load. Please check Space application logs."} try: raw_img = Image.fromarray(image).convert('RGB') inputs = processor(images=raw_img, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) probabilities = F.softmax(outputs.logits, dim=1)[0] target_layer = model.swin.encoder.layers[STAGE_INDEX] heatmap = get_gradcam(model, inputs["pixel_values"], target_layer) cam_pil = apply_overlay(raw_img, heatmap) top_prob, top_idx = torch.topk(probabilities, k=min(5, len(str_idx_to_class))) confidences = {str_idx_to_class.get(str(idx.item()), f"Class {idx.item()}"): float(prob.item()) for prob, idx in zip(top_prob, top_idx)} return cam_pil, confidences except Exception as e: return None, {"Error": str(e)} # --- Gradio Blocks Interface --- with gr.Blocks(title="Thai Snake Classifier", theme="soft") as demo: gr.Markdown("# Thai Snake Classifier") gr.Markdown("Classify Thai snakes with Swin-Base and Grad-CAM visualization.") with gr.Row(): with gr.Column(): input_img = gr.Image(label="Upload Snake Image", type="numpy") btn = gr.Button("Classify", variant="primary") with gr.Column(): output_cam = gr.Image(label="Grad-CAM Visualization", type="pil") output_labels = gr.Label(num_top_classes=5, label="Confidences") # Add Examples with explicit labels matching model id2label gr.Examples( examples=[ ["examples/Trimeresurus_albolabris_1.png", "Trimeresurus_albolabris"], ["examples/Trimeresurus_albolabris_2.png", "Trimeresurus_albolabris"], ["examples/Trimeresurus_albolabris_3.png", "Trimeresurus_albolabris"], ["examples/Trimeresurus_macrops_1.png", "Trimeresurus_macrops"], ["examples/Trimeresurus_macrops_2.png", "Trimeresurus_macrops"], ["examples/Trimeresurus_macrops_3.png", "Trimeresurus_macrops"], ["examples/Ptyas_mucosa_1.png", "Ptyas_mucosa"], ["examples/Ptyas_mucosa_2.png", "Ptyas_mucosa"], ["examples/Oligodon_taeniatus_1.png", "Oligodon_taeniatus"], ["examples/Oligodon_taeniatus_2.png", "Oligodon_taeniatus"], ["examples/Oligodon_taeniatus_3.png", "Oligodon_taeniatus"] ], inputs=[input_img, gr.Textbox(visible=False, label="Species Label")], # Dummy input for label outputs=[output_cam, output_labels], fn=predict, cache_examples=False, ) btn.click(fn=predict, inputs=input_img, outputs=[output_cam, output_labels]) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860)