fishit / app.py
aryn25's picture
Create app.py
cb3c56d verified
from transformers import BlipProcessor, BlipForConditionalGeneration
from sentence_transformers import SentenceTransformer
import gradio as gr
from PIL import Image
import torch
import requests
from bs4 import BeautifulSoup
# Load BLIP Model
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
# Map common fish names to FishBase scientific names
name_map = {
"pufferfish": "Tetraodon",
"stonefish": "Synanceia",
"lionfish": "Pterois",
"tuna": "Thunnus",
"salmon": "Salmo-salar",
"catfish": "Ictalurus",
"tilapia": "Oreochromis"
}
# Poisonous species (scientific names)
poisonous_species = ["Tetraodon", "Synanceia", "Pterois"]
# FishBase scraping function
def get_fishbase_summary(scientific_name):
search_url = f"https://www.fishbase.se/summary/{scientific_name}.html"
try:
response = requests.get(search_url, timeout=10)
if response.status_code != 200:
return f"FishBase entry not found for: {scientific_name}"
soup = BeautifulSoup(response.text, "html.parser")
summary_section = soup.find("div", {"id": "ssummary"})
if summary_section:
paragraphs = summary_section.find_all("p")
text = "\n\n".join(p.get_text(strip=True) for p in paragraphs if p.get_text(strip=True))
return text or f"No summary available for {scientific_name}"
else:
return f"No detailed summary found for {scientific_name}"
except Exception as e:
return f"Error fetching FishBase data for {scientific_name}: {str(e)}"
# Fish identification function
def identify_fish(image):
# Step 1: Generate caption from image
inputs = blip_processor(image, return_tensors="pt")
out = blip_model.generate(**inputs)
caption = blip_processor.decode(out[0], skip_special_tokens=True)
# Step 2: Extract fish name from caption
fish_name = None
for name in name_map:
if name in caption.lower():
fish_name = name
break
if not fish_name:
return f"❌ Could not identify a known fish species in the image caption: '{caption}'"
# Step 3: Lookup in FishBase
scientific_name = name_map[fish_name]
summary = get_fishbase_summary(scientific_name)
# Step 4: Check toxicity
is_poisonous = "Yes 🧪" if scientific_name in poisonous_species else "No ✅"
# Step 5: Final Output
return f"**Image Caption:** {caption}\n\n**Detected Fish:** {fish_name.title()}\n**Scientific Name:** {scientific_name}\n**Poisonous:** {is_poisonous}\n\n**📚 FishBase Info:**\n{summary}"
# Gradio UI
demo = gr.Interface(
fn=identify_fish,
inputs=gr.Image(type="pil"),
outputs="markdown",
title="🐟 Smart Fish Identifier (BLIP + FishBase)",
description="Upload a fish image. We use BLIP to describe the fish, match it with known species, then fetch info from FishBase to check if it's poisonous."
)
if __name__ == '__main__':
demo.launch()