app.py

import streamlit as st
import torch
from PIL import Image
import open_clip

from transformers import (
    BlipProcessor,
    BlipForConditionalGeneration
)

st.set_page_config(page_title="Zero Shot Image Classification", layout="wide")

st.title("Zero Shot Image Classification")
st.write("BiomedCLIP + RemoteCLIP + AgriCLIP + BLIP")

device = "cpu"


# --------------------------------------------------
# LOAD MODELS
# --------------------------------------------------

@st.cache_resource
def load_models():

    # ---------- BIOMEDCLIP ----------
    biomed_model, _, biomed_preprocess = open_clip.create_model_and_transforms(
        "hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"
    )

    biomed_tokenizer = open_clip.get_tokenizer(
        "hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"
    )

    biomed_model = biomed_model.to(device).eval()


    # ---------- REMOTECLIP ----------
    remote_model, _, remote_preprocess = open_clip.create_model_and_transforms(
        "ViT-B-32",
        pretrained="laion2b_s34b_b79k"
    )

    remote_tokenizer = open_clip.get_tokenizer("ViT-B-32")

    remote_model = remote_model.to(device).eval()


    # ---------- AGRICLIP ----------
   # ---------- AGRICULTURE CLIP ----------
    agri_model, _, agri_preprocess = open_clip.create_model_and_transforms(
    "ViT-B-32",
    pretrained="laion2b_s34b_b79k"
)

    agri_tokenizer = open_clip.get_tokenizer("ViT-B-32")

    agri_model = agri_model.to(device).eval()


    # ---------- BLIP ----------
    blip_processor = BlipProcessor.from_pretrained(
        "Salesforce/blip-image-captioning-base"
    )

    blip_model = BlipForConditionalGeneration.from_pretrained(
        "Salesforce/blip-image-captioning-base"
    ).to(device).eval()


    return (
        biomed_model,
        biomed_preprocess,
        biomed_tokenizer,
        remote_model,
        remote_preprocess,
        remote_tokenizer,
        agri_model,
        agri_preprocess,
        agri_tokenizer,
        blip_processor,
        blip_model
    )


(
    biomed_model,
    biomed_preprocess,
    biomed_tokenizer,
    remote_model,
    remote_preprocess,
    remote_tokenizer,
    agri_model,
    agri_preprocess,
    agri_tokenizer,
    blip_processor,
    blip_model
) = load_models()


# --------------------------------------------------
# DATASET CLASSES
# --------------------------------------------------

DATASETS = {
    "medical": ["pneumonia", "Normal"],
    "skin_disease": ["Normal Skin", "eczema", "Melanoma", "psoriasis"],
    "satellite": ["HIGHWAY", "RIVER", "INDUSTRIAL", "FOREST", "CROP"],
    "agriculture": ["POWDERY_MILDEW", "HEALTHY", "RUST", "EARLY BLIGHT", "LATE BLIGHT"]
}


# --------------------------------------------------
# PROMPT TEMPLATES
# --------------------------------------------------

templates = {

"medical": {

"pneumonia":[
"a chest x ray showing pneumonia infection",
"a lung radiology scan with cloudy pneumonia regions",
"a chest radiograph indicating pneumonia"
],

"Normal":[
"a normal chest x ray with healthy lungs",
"a lung radiology scan without infection",
"a clear chest radiograph"
]

},

"skin_disease": {

"Normal Skin":[
"a dermatology image of healthy skin",
"a clinical photo of normal human skin"
],

"eczema":[
"a dermatology image showing eczema rash",
"a red irritated eczema patch on skin"
],

"Melanoma":[
"a dermatology image showing melanoma lesion",
"a dark irregular melanoma skin lesion"
],

"psoriasis":[
"a dermatology image showing psoriasis scales",
"a thick red psoriasis plaque"
]

},

"satellite": {

"HIGHWAY":[
"a satellite image showing a highway",
"an aerial view of highway road"
],

"RIVER":[
"a satellite image of a river",
"a remote sensing image of river channel"
],

"INDUSTRIAL":[
"a satellite image of an industrial area",
"a remote sensing image showing factories"
],

"FOREST":[
"a satellite image of dense forest",
"a remote sensing image showing forest canopy"
],

"CROP":[
"a satellite image of crop fields",
"a remote sensing image of farmland"
]

},

"agriculture": {

"POWDERY_MILDEW":[
"a plant leaf infected with powdery mildew",
"a leaf covered with white powder fungus"
],

"HEALTHY":[
"a healthy green plant leaf",
"a crop leaf without disease"
],

"RUST":[
"a crop leaf infected with rust disease",
"a plant leaf with rust spots"
],

"EARLY BLIGHT":[
"a plant leaf with early blight spots"
],

"LATE BLIGHT":[
"a plant leaf infected with late blight"
]

}

}


# --------------------------------------------------
# EXPLANATION BLOCK
# --------------------------------------------------

explanations = {

"medical": {
"pneumonia":
"""
The uploaded chest X-ray image appears to indicate pneumonia.  
Pneumonia is a lung infection that causes inflammation in the air sacs, which may fill with fluid or pus.  
Common visual indicators in X-rays include cloudy regions or opacities in the lung fields.  
Early detection is important because pneumonia can affect breathing and oxygen exchange.  
Medical professionals typically confirm the diagnosis using radiological evaluation and clinical symptoms.
""",

"Normal":
"""
The chest X-ray image appears to show healthy lungs with no visible infection.  
Normal lungs in radiographs usually show clear lung fields without abnormal opacities.  
The diaphragm and lung boundaries appear well defined in healthy scans.  
This suggests that the lungs are functioning normally without signs of pneumonia or infection.  
A medical professional would still review the image alongside patient symptoms for confirmation.
"""
},

"skin_disease": {

"eczema":
"""
The image appears to show signs consistent with eczema.  
Eczema is a skin condition that causes redness, inflammation, dryness, and itching.  
It often appears as irritated patches of skin with uneven texture.  
Environmental triggers, allergies, or immune responses may cause flare-ups.  
Dermatologists typically diagnose eczema through visual inspection and medical history.
""",

"Melanoma":
"""
The image may contain characteristics associated with melanoma.  
Melanoma is a serious form of skin cancer that develops in pigment-producing cells.  
Visual indicators can include dark, irregularly shaped lesions with uneven borders.  
Early detection is critical because melanoma can spread to other organs.  
A dermatologist should examine suspicious lesions and may perform a biopsy.
""",

"psoriasis":
"""
The image appears to show features consistent with psoriasis.  
Psoriasis is a chronic autoimmune condition that causes rapid skin cell buildup.  
This results in thick, red patches with silvery scales on the skin.  
The condition may cause itching or discomfort depending on severity.  
Medical treatment usually focuses on reducing inflammation and slowing skin cell growth.
""",

"Normal Skin":
"""
The image appears to show healthy skin without visible dermatological abnormalities.  
Normal skin typically has an even tone and smooth surface texture.  
There are no visible lesions, scales, or inflammatory patches in the image.  
This suggests that the skin does not currently show signs of common dermatological diseases.  
However, regular skin monitoring is still recommended for early detection of conditions.
"""
},

"satellite": {

"HIGHWAY":
"""
The satellite image appears to contain a highway or major roadway.  
Highways are large transportation corridors designed for high-speed vehicle movement.  
From satellite imagery, they appear as long linear structures crossing landscapes.  
They connect cities, industrial zones, and transportation hubs.  
Remote sensing is commonly used to monitor infrastructure development and urban planning.
""",

"RIVER":
"""
The satellite image likely shows a river or flowing water body.  
Rivers appear as long winding structures that transport water across landscapes.  
They are important for agriculture, ecosystems, and human settlements.  
Satellite imagery is often used to monitor river flow, flooding, and environmental changes.  
Remote sensing helps researchers track water resources over time.
""",

"FOREST":
"""
The satellite image appears to contain dense forest vegetation.  
Forests typically appear as large green areas with textured canopy patterns in aerial imagery.  
They play a critical role in maintaining biodiversity and regulating climate.  
Satellite monitoring is widely used to detect deforestation and vegetation changes.  
Remote sensing helps scientists study environmental conservation.
"""
},

"agriculture": {

"POWDERY_MILDEW":
"""
The plant leaf in the image appears to show powdery mildew infection.  
Powdery mildew is a fungal disease that creates white powder-like patches on leaves.  
It can reduce photosynthesis and weaken plant health if untreated.  
This disease spreads quickly in humid environments with limited airflow.  
Farmers usually manage it using fungicides and improved crop management practices.
""",

"RUST":
"""
The plant leaf shows symptoms consistent with rust disease.  
Rust infections create orange or brown spots on the leaf surface.  
This fungal disease spreads through airborne spores.  
Severe infections can reduce crop yield and plant health.  
Agricultural monitoring helps detect plant diseases early for better crop protection.
""",

"EARLY BLIGHT":
"""
The leaf may show symptoms of early blight disease.  
Early blight is a fungal infection commonly affecting tomato and potato plants.  
It often produces circular brown spots with concentric rings on leaves.  
If left untreated, the disease may spread across the plant.  
Proper crop rotation and fungicide treatment help control early blight.
""",

"LATE BLIGHT":
"""
The leaf may be affected by late blight disease.  
Late blight is a severe plant disease that affects crops such as potatoes and tomatoes.  
Symptoms often include dark lesions and rapid leaf decay.  
The disease spreads quickly in cool and humid conditions.  
Early detection helps farmers take preventive measures to protect crops.
""",

"HEALTHY":
"""
The plant leaf appears healthy with no visible disease symptoms.  
Healthy leaves typically have a uniform green color and smooth surface.  
There are no visible fungal spots, discoloration, or lesions.  
This indicates that the plant is likely growing under good conditions.  
Regular monitoring helps maintain crop health and productivity.
"""
}

}


# --------------------------------------------------
# SIDEBAR
# --------------------------------------------------

dataset_key = st.sidebar.selectbox(
    "Select Dataset",
    list(DATASETS.keys())
)


# --------------------------------------------------
# IMAGE UPLOAD
# --------------------------------------------------

uploaded_file = st.file_uploader(
    "Upload Image",
    type=["jpg","jpeg","png"]
)


if uploaded_file:

    image = Image.open(uploaded_file).convert("RGB")

    st.image(image, width=400)

    labels = DATASETS[dataset_key]

    text_queries = []
    label_mapping = []

    for label in labels:
        for caption in templates[dataset_key][label]:
            text_queries.append(caption)
            label_mapping.append(label)


    # --------------------------------------------------
    # MODEL SELECTION
    # --------------------------------------------------

    if dataset_key in ["medical","skin_disease"]:

        img = biomed_preprocess(image).unsqueeze(0).to(device)
        text = biomed_tokenizer(text_queries)

        with torch.no_grad():
            image_features = biomed_model.encode_image(img)
            text_features = biomed_model.encode_text(text)

        similarity = (image_features @ text_features.T).softmax(dim=-1)


    elif dataset_key == "satellite":

        img = remote_preprocess(image).unsqueeze(0).to(device)
        text = remote_tokenizer(text_queries)

        with torch.no_grad():
            image_features = remote_model.encode_image(img)
            text_features = remote_model.encode_text(text)

        similarity = (image_features @ text_features.T).softmax(dim=-1)


    elif dataset_key == "agriculture":

        img = agri_preprocess(image).unsqueeze(0).to(device)
        text = agri_tokenizer(text_queries)

        with torch.no_grad():
            image_features = agri_model.encode_image(img)
            text_features = agri_model.encode_text(text)

        similarity = (image_features @ text_features.T).softmax(dim=-1)


    conf, idx = torch.max(similarity, dim=1)

    predicted_class = label_mapping[idx.item()]

    st.success(f"Prediction: {predicted_class}")
    st.metric("Confidence", f"{conf.item():.2%}")


    # --------------------------------------------------
    # BLIP IMAGE CAPTION
    # --------------------------------------------------

    blip_inputs = blip_processor(images=image, return_tensors="pt").to(device)

    with torch.no_grad():
        caption_ids = blip_model.generate(**blip_inputs)

    caption = blip_processor.decode(caption_ids[0], skip_special_tokens=True)

    st.subheader("Image Description (BLIP)")
    st.write(caption)


    # --------------------------------------------------
    # SHOW EXPLANATION
    # --------------------------------------------------

    st.subheader("Explanation")

    if dataset_key in explanations and predicted_class in explanations[dataset_key]:
        st.write(explanations[dataset_key][predicted_class])
    else:
        st.write("The prediction was generated by comparing the uploaded image with multiple textual descriptions using a zero-shot vision-language model.")