Update app.py

app.py

@@ -2,36 +2,39 @@ import gradio as gr
 import torch
 import torchvision.transforms as transforms
 from torchvision import models
+from torchvision.models import ResNet18_Weights
 from PIL import Image
 import matplotlib.pyplot as plt
 import numpy as np
 import cv2
-import pandas as pd
-from datetime import datetime
 from gtts import gTTS
 import wikipedia
+import json
+import urllib.request
 import os
 
 # ===============================
-# LOAD MODEL
+# LOAD MODEL (Updated for Torchvision 0.13+)
 # ===============================
-model = models.resnet18(pretrained=True)
+weights = ResNet18_Weights.DEFAULT
+model = models.resnet18(weights=weights)
 model.eval()
 
 # ===============================
 # LABELS
 # ===============================
-import json, urllib.request
 url = "https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json"
-class_idx = json.load(urllib.request.urlopen(url))
+with urllib.request.urlopen(url) as response:
+    class_idx = json.load(response)
 labels = [class_idx[str(k)][1] for k in range(len(class_idx))]
 
 # ===============================
 # TRANSFORM
 # ===============================
 transform = transforms.Compose([
-    transforms.Resize((224,224)),
-    transforms.ToTensor()
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
 # ===============================
@@ -39,135 +42,176 @@ transform = transforms.Compose([
 # ===============================
 def get_wiki_info(query):
     try:
+        # Wikipedia sometimes needs specific search terms
         return wikipedia.summary(query, sentences=2)
     except:
-        return "No info available."
+        return "No additional detailed information available on Wikipedia."
+
+# ===============================
+# NATURAL VOICE
+# ===============================
+def generate_speech(label, wiki_info, lang):
+    if lang == "Telugu":
+        text = f"ఈ చిత్రం {label} ను చూపిస్తుంది. {wiki_info}"
+        tts = gTTS(text=text, lang='te')
+    elif lang == "Hindi":
+        text = f"यह चित्र {label} दिखाता है। {wiki_info}"
+        tts = gTTS(text=text, lang='hi')
+    else:
+        text = f"This image appears to show a {label}. {wiki_info}"
+        tts = gTTS(text=text, lang='en')
+
+    audio_file = "voice.mp3"
+    tts.save(audio_file)
+    return audio_file
 
 # ===============================
 # GRAD-CAM
 # ===============================
 def generate_gradcam(image):
-    img = transform(image).unsqueeze(0)
-    img.requires_grad = True
-
+    # Preprocess
+    img_tensor = transform(image).unsqueeze(0)
+    
     features = []
-    def hook_fn(module, input, output):
+    gradients = []
+
+    def forward_hook(module, input, output):
         features.append(output)
 
-    handle = model.layer4.register_forward_hook(hook_fn)
+    def backward_hook(module, grad_in, grad_out):
+        gradients.append(grad_out[0])
 
-    output = model(img)
+    # Register hooks on the last convolutional layer
+    handle_f = model.layer4.register_forward_hook(forward_hook)
+    handle_b = model.layer4.register_full_backward_hook(backward_hook)
+
+    # Forward pass
+    output = model(img_tensor)
     pred_class = output.argmax()
 
+    # Backward pass
     model.zero_grad()
     output[0, pred_class].backward()
 
-    gradients = model.layer4[1].conv2.weight.grad
-
-    pooled_gradients = torch.mean(gradients, dim=[0,2,3])
-
-    fmap = features[0].squeeze()
+    # Process CAM
+    fmap = features[0].detach()
+    grad = gradients[0].detach()
+    weights_cam = torch.mean(grad, dim=[2, 3])[0]
+    cam = torch.zeros(fmap.shape[2:], dtype=torch.float32)
 
-    for i in range(len(pooled_gradients)):
-        fmap[i,:,:] *= pooled_gradients[i]
+    for i, w in enumerate(weights_cam):
+        cam += w * fmap[0, i, :, :]
 
-    heatmap = torch.mean(fmap, dim=0).detach().numpy()
-    heatmap = np.maximum(heatmap, 0)
-    heatmap /= heatmap.max()
+    cam = np.maximum(cam.numpy(), 0)
+    if cam.max() > 0:
+        cam = cam / cam.max()
 
-    heatmap = cv2.resize(heatmap, (224,224))
-    heatmap = np.uint8(255 * heatmap)
+    cam = cv2.resize(cam, (image.size[0], image.size[1]))
+    heatmap = np.uint8(255 * cam)
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
 
-    img_np = np.array(image.resize((224,224)))
+    img_np = np.array(image)
     superimposed = heatmap * 0.4 + img_np
 
-    handle.remove()
+    # Cleanup hooks
+    handle_f.remove()
+    handle_b.remove()
 
     return superimposed.astype(np.uint8)
 
 # ===============================
 # PREDICT
 # ===============================
-def predict(image):
-    img = transform(image).unsqueeze(0)
+def predict(image, lang):
+    if image is None: return None, None, None, None, None
+    
+    img_tensor = transform(image).unsqueeze(0)
 
     with torch.no_grad():
-        outputs = model(img)
+        outputs = model(img_tensor)
         probs = torch.nn.functional.softmax(outputs[0], dim=0)
 
     top3 = torch.topk(probs, 3)
-
     results = []
     for i in range(3):
         label = labels[top3.indices[i]]
-        conf = float(top3.values[i])*100
+        conf = float(top3.values[i]) * 100
         results.append((label, conf))
 
     top_label = results[0][0]
-
-    # Wikipedia
     wiki = get_wiki_info(top_label)
-
-    # Grad-CAM
+    
+    # Temporarily enable grads for Grad-CAM
+    model.train(False) 
     heatmap = generate_gradcam(image)
 
     # Graph
     names = [r[0] for r in results]
     values = [r[1] for r in results]
-    plt.figure()
-    plt.bar(names, values)
+    fig = plt.figure(figsize=(8, 4))
+    plt.bar(names, values, color='skyblue')
+    plt.title("Top 3 Predictions")
+    plt.ylabel("Confidence %")
 
-    # Voice
-    text = f"{top_label}. {wiki}"
-    tts = gTTS(text=text)
-    tts.save("voice.mp3")
+    audio = generate_speech(top_label, wiki, lang)
+    text_output = "\n".join([f"{r[0]} : {r[1]:.2f}%" for r in results])
 
-    return results, wiki, heatmap, plt, "voice.mp3"
+    return text_output, wiki, heatmap, fig, audio
 
 # ===============================
-# REAL-TIME VIDEO
+# REAL-TIME WEBCAM
 # ===============================
 def webcam_predict(frame):
+    if frame is None: return None
     image = Image.fromarray(frame)
-    res, _, _, _, _ = predict(image)
+    img_tensor = transform(image).unsqueeze(0)
 
-    label = res[0][0]
-    cv2.putText(frame, label, (20,40),
-                cv2.FONT_HERSHEY_SIMPLEX, 1,
-                (0,255,0), 2)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = torch.nn.functional.softmax(outputs[0], dim=0)
 
+    label = labels[probs.argmax()]
+    cv2.putText(frame, f"Detected: {label}", (20, 50),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
     return frame
 
 # ===============================
-# PREMIUM UI
+# PREMIUM UI (Gradio 6.0 Compatibility)
 # ===============================
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-
+with gr.Blocks(theme="soft") as demo:
     gr.Markdown("# 🔥 Next-Gen AI Image Recognition System")
-    gr.Markdown("### With Grad-CAM + Real-Time Detection + Wikipedia + Voice")
-
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Upload Image")
-        output_heatmap = gr.Image(label="Explainable AI (Grad-CAM)")
-
-    predict_btn = gr.Button("🔍 Analyze Image")
-
-    output_text = gr.Textbox(label="Predictions")
-    output_wiki = gr.Textbox(label="Wikipedia Info")
-    output_graph = gr.Plot(label="Confidence Graph")
-    output_audio = gr.Audio(label="Voice Output")
-
-    predict_btn.click(
-        fn=predict,
-        inputs=image_input,
-        outputs=[output_text, output_wiki, output_heatmap, output_graph, output_audio]
+    gr.Markdown("### Grad-CAM + Real-Time Detection + Wikipedia + Natural Voice")
+
+    with gr.Tab("Static Analysis"):
+        with gr.Row():
+            image_input = gr.Image(type="pil", label="Upload Image", sources=["upload", "clipboard"])
+            heatmap_output = gr.Image(label="Explainable AI (Grad-CAM)")
+
+        lang = gr.Dropdown(["English", "Telugu", "Hindi"], label="Voice Language", value="English")
+        btn = gr.Button("🔍 Analyze Image", variant="primary")
+
+        with gr.Row():
+            output_text = gr.Textbox(label="Predictions")
+            output_wiki = gr.Textbox(label="Wikipedia Info")
+        
+        with gr.Row():
+            output_graph = gr.Plot(label="Confidence Graph")
+            output_audio = gr.Audio(label="Voice Explanation", autoplay=True)
+
+    with gr.Tab("Real-Time Webcam"):
+        gr.Markdown("## 🎥 Live Stream Detection")
+        # Gradio 6.x Streaming Setup
+        with gr.Row():
+            webcam_input = gr.Image(sources=["webcam"], streaming=True, label="Webcam Input")
+            webcam_output = gr.Image(label="Processed Feed")
+        
+        webcam_input.stream(webcam_predict, inputs=[webcam_input], outputs=[webcam_output])
+
+    btn.click(
+        predict,
+        inputs=[image_input, lang],
+        outputs=[output_text, output_wiki, heatmap_output, output_graph, output_audio]
     )
 
-    gr.Markdown("## 🎥 Real-Time Webcam Detection")
-    webcam = gr.Image(source="webcam", streaming=True)
-    webcam_output = gr.Image()
-
-    webcam.stream(webcam_predict, inputs=webcam, outputs=webcam_output)
-
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()