Spaces:

hamzabinabdulmajeed
/

Archaeology_Project

Sleeping

Create app.py

533cf2c verified 5 months ago

6.02 kB

import os
import time
import torch
import torch.nn as nn
from PIL import Image
import pandas as pd
import gradio as gr
from torchvision import models, transforms
from torch.autograd import Function

# 1. Dataset Configuration (Matching your ArchaeoMind settings)
CLASS_NAMES = ['artifact', 'non_artifact']
NUM_CLASSES = len(CLASS_NAMES)

# 2. Define Network Architectures to Load Weights Safely

# --- Model A: ResNet-18 ---
resnet18 = models.resnet18(weights=None)
num_features_res = resnet18.fc.in_features
resnet18.fc = nn.Sequential(
    nn.Linear(num_features_res, 256),
    nn.ReLU(inplace=True),
    nn.Dropout(p=0.4),
    nn.Linear(256, NUM_CLASSES)
)

# --- Model B: VGG16 Baseline ---
vgg16_baseline = models.vgg16(weights=None)
num_features_vgg = vgg16_baseline.classifier[0].in_features
vgg16_baseline.classifier = nn.Sequential(
    nn.Linear(num_features_vgg, 256),
    nn.ReLU(inplace=True),
    nn.Dropout(p=0.5),
    nn.Linear(256, NUM_CLASSES)
)

# --- Model C: VGG16 Domain-Adversarial (DANN) ---
class GradientReversalLayer(Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.alpha = alpha
        return x.view_as(x)
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output.neg() * ctx.alpha, None

def grad_reverse(x, alpha=1.0):
    return GradientReversalLayer.apply(x, alpha)

class VGG16_DANN(nn.Module):
    def __init__(self, num_classes=2):
        super(VGG16_DANN, self).__init__()
        vgg_base = models.vgg16(weights=None)
        self.feature_extractor = vgg_base.features
        self.pool = nn.AdaptiveAvgPool2d((4, 4))
        in_features = 512 * 4 * 4

        self.class_classifier = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(256, num_classes)
        )
        self.domain_classifier = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(256, 2)
        )

    def forward(self, x, alpha=1.0):
        features = self.feature_extractor(x)
        features = self.pool(features)
        features = torch.flatten(features, 1)
        class_output = self.class_classifier(features)
        reversed_features = grad_reverse(features, alpha)
        domain_output = self.domain_classifier(reversed_features)
        return class_output, domain_output

vgg16_dann = VGG16_DANN(num_classes=NUM_CLASSES)

# 3. Load Model Checkpoints onto CPU
models_dict = {
    "ResNet-18": resnet18,
    "VGG16 Baseline": vgg16_baseline,
    "VGG16 DANN (Domain Adaptation)": vgg16_dann
}

files_dict = {
    "ResNet-18": "resnet18_artifacts.pth",
    "VGG16 Baseline": "vgg16_artifacts_baseline.pth",
    "VGG16 DANN (Domain Adaptation)": "vgg16_dann_artifacts.pth"
}

for name, model in models_dict.items():
    file_path = files_dict[name]
    if os.path.exists(file_path):
        model.load_state_dict(torch.load(file_path, map_location=torch.device('cpu')))
        model.eval()
        print(f"Successfully loaded {name} weights structure.")
    else:
        print(f"WARNING: Checkpoint file '{file_path}' missing!")

# 4. Input Transforms Pipeline (Matches your 128x128 3-Channel Grayscale specs)
img_transforms = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 5. Inference Benchmarking Engine
def benchmark_inference(input_image):
    if input_image is None:
        return None
        
    img_tensor = img_transforms(input_image).unsqueeze(0)
    results_data = []
    
    with torch.no_grad():
        for name, model in models_dict.items():
            # Warm-up pass
            if name == "VGG16 DANN (Domain Adaptation)":
                _, _ = model(img_tensor, alpha=0.0)
            else:
                _ = model(img_tensor)
            
            # Performance timing run
            start_time = time.perf_counter()
            if name == "VGG16 DANN (Domain Adaptation)":
                outputs, _ = model(img_tensor, alpha=0.0)  # Ignore domain output head here
            else:
                outputs = model(img_tensor)
                
            probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
            best_prob, best_class_idx = torch.max(probabilities, dim=0)
            end_time = time.perf_counter()
            
            latency_ms = (end_time - start_time) * 1000
            predicted_label = CLASS_NAMES[best_class_idx.item()]
            confidence_pct = f"{best_prob.item() * 100:.2f}%"
            
            results_data.append({
                "Model Architecture": name,
                "Classification Prediction": predicted_label,
                "Confidence Score": confidence_pct,
                "Inference Latency": f"{latency_ms:.2f} ms"
            })
            
    return pd.DataFrame(results_data)

# 6. Gradio Interface Layout
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔍 ArchaeoMind: Artifact Detection & Speed Analysis Space")
    gr.Markdown("Upload an excavation find image to classify it as an artifact or non-artifact while profiling model efficiency metrics across architectures.")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_img_slot = gr.Image(type="pil", label="Excavation Find Photo")
            submit_btn = gr.Button("Analyze Image", variant="primary")
            
        with gr.Column(scale=2):
            output_table_slot = gr.Dataframe(
                headers=["Model Architecture", "Classification Prediction", "Confidence Score", "Inference Latency"],
                datatype=["str", "str", "str", "str"],
                label="Benchmark Results Table"
            )
            
    submit_btn.click(
        fn=benchmark_inference, 
        inputs=input_img_slot, 
        outputs=output_table_slot
    )

demo.launch()

	import os
	import time
	import torch
	import torch.nn as nn
	from PIL import Image
	import pandas as pd
	import gradio as gr
	from torchvision import models, transforms
	from torch.autograd import Function

	# 1. Dataset Configuration (Matching your ArchaeoMind settings)
	CLASS_NAMES = ['artifact', 'non_artifact']
	NUM_CLASSES = len(CLASS_NAMES)

	# 2. Define Network Architectures to Load Weights Safely

	# --- Model A: ResNet-18 ---
	resnet18 = models.resnet18(weights=None)
	num_features_res = resnet18.fc.in_features
	resnet18.fc = nn.Sequential(
	nn.Linear(num_features_res, 256),
	nn.ReLU(inplace=True),
	nn.Dropout(p=0.4),
	nn.Linear(256, NUM_CLASSES)
	)

	# --- Model B: VGG16 Baseline ---
	vgg16_baseline = models.vgg16(weights=None)
	num_features_vgg = vgg16_baseline.classifier[0].in_features
	vgg16_baseline.classifier = nn.Sequential(
	nn.Linear(num_features_vgg, 256),
	nn.ReLU(inplace=True),
	nn.Dropout(p=0.5),
	nn.Linear(256, NUM_CLASSES)
	)

	# --- Model C: VGG16 Domain-Adversarial (DANN) ---
	class GradientReversalLayer(Function):
	@staticmethod
	def forward(ctx, x, alpha):
	ctx.alpha = alpha
	return x.view_as(x)
	@staticmethod
	def backward(ctx, grad_output):
	return grad_output.neg() * ctx.alpha, None

	def grad_reverse(x, alpha=1.0):
	return GradientReversalLayer.apply(x, alpha)

	class VGG16_DANN(nn.Module):
	def __init__(self, num_classes=2):
	super(VGG16_DANN, self).__init__()
	vgg_base = models.vgg16(weights=None)
	self.feature_extractor = vgg_base.features
	self.pool = nn.AdaptiveAvgPool2d((4, 4))
	in_features = 512 * 4 * 4

	self.class_classifier = nn.Sequential(
	nn.Linear(in_features, 256),
	nn.ReLU(inplace=True),
	nn.Dropout(p=0.5),
	nn.Linear(256, num_classes)
	)
	self.domain_classifier = nn.Sequential(
	nn.Linear(in_features, 256),
	nn.ReLU(inplace=True),
	nn.Dropout(p=0.5),
	nn.Linear(256, 2)
	)

	def forward(self, x, alpha=1.0):
	features = self.feature_extractor(x)
	features = self.pool(features)
	features = torch.flatten(features, 1)
	class_output = self.class_classifier(features)
	reversed_features = grad_reverse(features, alpha)
	domain_output = self.domain_classifier(reversed_features)
	return class_output, domain_output

	vgg16_dann = VGG16_DANN(num_classes=NUM_CLASSES)

	# 3. Load Model Checkpoints onto CPU
	models_dict = {
	"ResNet-18": resnet18,
	"VGG16 Baseline": vgg16_baseline,
	"VGG16 DANN (Domain Adaptation)": vgg16_dann
	}

	files_dict = {
	"ResNet-18": "resnet18_artifacts.pth",
	"VGG16 Baseline": "vgg16_artifacts_baseline.pth",
	"VGG16 DANN (Domain Adaptation)": "vgg16_dann_artifacts.pth"
	}

	for name, model in models_dict.items():
	file_path = files_dict[name]
	if os.path.exists(file_path):
	model.load_state_dict(torch.load(file_path, map_location=torch.device('cpu')))
	model.eval()
	print(f"Successfully loaded {name} weights structure.")
	else:
	print(f"WARNING: Checkpoint file '{file_path}' missing!")

	# 4. Input Transforms Pipeline (Matches your 128x128 3-Channel Grayscale specs)
	img_transforms = transforms.Compose([
	transforms.Resize((128, 128)),
	transforms.Grayscale(num_output_channels=3),
	transforms.ToTensor(),
	transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
	])

	# 5. Inference Benchmarking Engine
	def benchmark_inference(input_image):
	if input_image is None:
	return None

	img_tensor = img_transforms(input_image).unsqueeze(0)
	results_data = []

	with torch.no_grad():
	for name, model in models_dict.items():
	# Warm-up pass
	if name == "VGG16 DANN (Domain Adaptation)":
	_, _ = model(img_tensor, alpha=0.0)
	else:
	_ = model(img_tensor)

	# Performance timing run
	start_time = time.perf_counter()
	if name == "VGG16 DANN (Domain Adaptation)":
	outputs, _ = model(img_tensor, alpha=0.0) # Ignore domain output head here
	else:
	outputs = model(img_tensor)

	probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
	best_prob, best_class_idx = torch.max(probabilities, dim=0)
	end_time = time.perf_counter()

	latency_ms = (end_time - start_time) * 1000
	predicted_label = CLASS_NAMES[best_class_idx.item()]
	confidence_pct = f"{best_prob.item() * 100:.2f}%"

	results_data.append({
	"Model Architecture": name,
	"Classification Prediction": predicted_label,
	"Confidence Score": confidence_pct,
	"Inference Latency": f"{latency_ms:.2f} ms"
	})

	return pd.DataFrame(results_data)

	# 6. Gradio Interface Layout
	with gr.Blocks(theme=gr.themes.Soft()) as demo:
	gr.Markdown("# 🔍 ArchaeoMind: Artifact Detection & Speed Analysis Space")
	gr.Markdown("Upload an excavation find image to classify it as an artifact or non-artifact while profiling model efficiency metrics across architectures.")

	with gr.Row():
	with gr.Column(scale=1):
	input_img_slot = gr.Image(type="pil", label="Excavation Find Photo")
	submit_btn = gr.Button("Analyze Image", variant="primary")

	with gr.Column(scale=2):
	output_table_slot = gr.Dataframe(
	headers=["Model Architecture", "Classification Prediction", "Confidence Score", "Inference Latency"],
	datatype=["str", "str", "str", "str"],
	label="Benchmark Results Table"
	)

	submit_btn.click(
	fn=benchmark_inference,
	inputs=input_img_slot,
	outputs=output_table_slot
	)

	demo.launch()