Update app.py

app.py

@@ -1,26 +1,45 @@
-import gradio as gr
+import streamlit as st
 import torch
 from PIL import Image
 from transformers import CLIPModel, CLIPProcessor
 from transformers import BlipProcessor, BlipForConditionalGeneration
-from datasets import load_dataset
+
+st.set_page_config(
+    page_title="AI Image Diagnostic System",
+    layout="wide"
+)
+
+st.title("🔬 AI Image Diagnostic System")
+st.write("CLIP + BLIP based AI diagnostic platform")
 
 device = "cpu"
 
-print("Loading models...")
+# Load models once
+@st.cache_resource
+def load_models():
 
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    clip_model = CLIPModel.from_pretrained(
+        "openai/clip-vit-base-patch32"
+    )
 
-blip_processor = BlipProcessor.from_pretrained(
-    "Salesforce/blip-image-captioning-base"
-)
+    clip_processor = CLIPProcessor.from_pretrained(
+        "openai/clip-vit-base-patch32"
+    )
+
+    blip_processor = BlipProcessor.from_pretrained(
+        "Salesforce/blip-image-captioning-base"
+    )
+
+    blip_model = BlipForConditionalGeneration.from_pretrained(
+        "Salesforce/blip-image-captioning-base"
+    )
+
+    return clip_model, clip_processor, blip_processor, blip_model
+
+
+clip_model, clip_processor, blip_processor, blip_model = load_models()
 
-blip_model = BlipForConditionalGeneration.from_pretrained(
-    "Salesforce/blip-image-captioning-base"
-)
 
-# DATASET LABELS
 DATASETS = {
     "medical": ["pneumonia", "Normal"],
     "skin_cancer": ["Normal Skin", "eczema", "Melanoma", "psoriasis"],
@@ -35,11 +54,33 @@ templates = {
     "agriculture": "a close-up leaf showing signs of {}"
 }
 
-def analyze(image, dataset):
 
-    labels = DATASETS[dataset]
+st.sidebar.header("Settings")
+
+dataset_key = st.sidebar.selectbox(
+    "Select Dataset Type",
+    list(DATASETS.keys())
+)
+
+uploaded_file = st.file_uploader(
+    "Upload Image",
+    type=["jpg", "jpeg", "png"]
+)
 
-    text_queries = [templates[dataset].format(l) for l in labels]
+if uploaded_file:
+
+    image = Image.open(uploaded_file).convert("RGB")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    labels = DATASETS[dataset_key]
+
+    text_queries = [
+        templates[dataset_key].format(l) for l in labels
+    ]
 
     inputs = clip_processor(
         text=text_queries,
@@ -49,40 +90,38 @@ def analyze(image, dataset):
     )
 
     with torch.no_grad():
+
         outputs = clip_model(**inputs)
+
         probs = outputs.logits_per_image.softmax(dim=1)
 
     conf, idx = torch.max(probs, dim=1)
 
-    detected_class = labels[idx.item()]
+    predicted_class = labels[idx.item()]
 
-    # BLIP caption generation
-    blip_inputs = blip_processor(images=image, return_tensors="pt")
+    with col2:
 
-    with torch.no_grad():
-        ids = blip_model.generate(**blip_inputs)
+        st.success(f"Prediction: {predicted_class}")
 
-    caption = blip_processor.decode(ids[0], skip_special_tokens=True)
+        st.metric(
+            label="Confidence",
+            value=f"{conf.item():.2%}"
+        )
 
-    return detected_class, float(conf), caption
+    blip_inputs = blip_processor(
+        images=image,
+        return_tensors="pt"
+    )
 
+    with torch.no_grad():
 
-interface = gr.Interface(
-    fn=analyze,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Dropdown(
-            choices=list(DATASETS.keys()),
-            label="Dataset Type"
-        )
-    ],
-    outputs=[
-        gr.Text(label="Predicted Class"),
-        gr.Number(label="Confidence"),
-        gr.Textbox(label="Description")
-    ],
-    title="AI Image Diagnostic System",
-    description="CLIP + BLIP based AI diagnostic model"
-)
+        caption_ids = blip_model.generate(**blip_inputs)
+
+    caption = blip_processor.decode(
+        caption_ids[0],
+        skip_special_tokens=True
+    )
+
+    st.subheader("Generated Description")
 
-interface.launch()
+    st.write(caption)