import json import os import gradio as gr import numpy as np import tensorflow as tf import torch import torch.nn as nn from PIL import Image from torchvision import transforms # ----------------------------- # Config # ----------------------------- PT_MODEL_PATH = "fatima_model.pth" TF_MODEL_PATH = "fatima_model.keras" META_PATH = "fatima_meta.json" DEFAULT_CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"] DEFAULT_IMAGE_SIZE = 150 DEFAULT_MEAN = [0.485, 0.456, 0.406] DEFAULT_STD = [0.229, 0.224, 0.225] device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def load_meta(): if os.path.exists(META_PATH): with open(META_PATH, "r", encoding="utf-8") as f: meta = json.load(f) class_names = meta.get("class_names", DEFAULT_CLASS_NAMES) image_size = int(meta.get("image_size", DEFAULT_IMAGE_SIZE)) mean = meta.get("imagenet_mean", DEFAULT_MEAN) std = meta.get("imagenet_std", DEFAULT_STD) return class_names, image_size, mean, std return DEFAULT_CLASS_NAMES, DEFAULT_IMAGE_SIZE, DEFAULT_MEAN, DEFAULT_STD CLASS_NAMES, IMAGE_SIZE, IMAGENET_MEAN, IMAGENET_STD = load_meta() class TorchCNN(nn.Module): def __init__(self, num_classes=6): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(128 * 18 * 18, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, num_classes), ) def forward(self, x): return self.classifier(self.features(x)) def load_pytorch_model(): ckpt = torch.load(PT_MODEL_PATH, map_location=device) class_names = ckpt.get("class_names", CLASS_NAMES) image_size = int(ckpt.get("image_size", IMAGE_SIZE)) model = TorchCNN(num_classes=len(class_names)) model.load_state_dict(ckpt["model_state"]) model.to(device).eval() return model, class_names, image_size def preprocess_pytorch(image: Image.Image, image_size: int): tfm = transforms.Compose( [ transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ] ) x = tfm(image.convert("RGB")).unsqueeze(0) return x.to(device) def load_tensorflow_model(): return tf.keras.models.load_model(TF_MODEL_PATH) def preprocess_tensorflow(image: Image.Image, image_size: int): image = image.convert("RGB").resize((image_size, image_size)) x = np.asarray(image, dtype=np.float32) / 255.0 mean = np.array(IMAGENET_MEAN, dtype=np.float32) std = np.array(IMAGENET_STD, dtype=np.float32) x = (x - mean) / std return np.expand_dims(x, axis=0) pt_model = None pt_class_names = None pt_image_size = None tf_model = None def predict(model_choice, image): global pt_model, pt_class_names, pt_image_size, tf_model if image is None: return "Please upload an image.", "", {} try: pil_img = image if isinstance(image, Image.Image) else Image.fromarray(image) if model_choice == "PyTorch": if pt_model is None: pt_model, pt_class_names, pt_image_size = load_pytorch_model() x = preprocess_pytorch(pil_img, pt_image_size) with torch.no_grad(): logits = pt_model(x) probs = torch.softmax(logits, dim=1).cpu().numpy()[0] pred_idx = int(np.argmax(probs)) label = pt_class_names[pred_idx] confidence = float(probs[pred_idx]) details = {pt_class_names[i]: float(probs[i]) for i in range(len(pt_class_names))} return f"Prediction: {label}", f"Confidence: {confidence:.2%}", details if tf_model is None: tf_model = load_tensorflow_model() x = preprocess_tensorflow(pil_img, IMAGE_SIZE) probs = tf_model.predict(x, verbose=0)[0] pred_idx = int(np.argmax(probs)) label = CLASS_NAMES[pred_idx] confidence = float(probs[pred_idx]) details = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))} return f"Prediction: {label}", f"Confidence: {confidence:.2%}", details except Exception as exc: return f"Inference error: {exc}", "", {} CUSTOM_CSS = """ .gradio-container { max-width: 1050px !important; } .main-card { border-radius: 20px; padding: 18px; background: linear-gradient(135deg, #0f172a 0%, #1e3a8a 45%, #1d4ed8 100%); color: white; } .main-title { font-size: 30px; font-weight: 800; margin-bottom: 6px; } .subtitle { color: #dbeafe; font-size: 14px; } .badge { display: inline-block; padding: 6px 10px; margin-right: 8px; border-radius: 999px; background: rgba(255,255,255,0.18); font-size: 12px; } """ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS, title="Intel Classifier") as demo: gr.HTML( """
Intel Image Classification
Choose a model, upload an image, and get the predicted class.
PyTorch + TensorFlow 6 Classes Image Size: 150x150
""" ) with gr.Row(): with gr.Column(scale=1): model_choice = gr.Dropdown( choices=["PyTorch", "TensorFlow"], value="PyTorch", label="Model", ) image_input = gr.Image(type="pil", label="Upload image") with gr.Row(): predict_btn = gr.Button("Predict", variant="primary") clear_btn = gr.Button("Clear") with gr.Column(scale=1): pred_text = gr.Textbox(label="Predicted class") conf_text = gr.Textbox(label="Confidence") probs = gr.Label(label="Class probabilities", num_top_classes=6) predict_btn.click( fn=predict, inputs=[model_choice, image_input], outputs=[pred_text, conf_text, probs], ) clear_btn.click( fn=lambda: ("", "", None, None), inputs=[], outputs=[pred_text, conf_text, probs, image_input], ) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860)