ossaili/27_Architectural_Styles_Classifier

1import sys2import PIL3import cv24import torch5import torchvision6import torch.nn as nn7from utils.save_load import load_model8import gradio as gr9from PIL import Image10from torchvision import transforms11import gradio as gr12from pytorch_grad_cam import GradCAM, AblationCAM, FullGrad, EigenGradCAM, LayerCAM13from pytorch_grad_cam.utils.image import show_cam_on_image14from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget15from pytorch_grad_cam import DeepFeatureFactorization16from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image, deprocess_image17import numpy as np18from typing import List19from matplotlib import pyplot as plt20from matplotlib.lines import Line2D21 22labels = [23    "Achaemenid architecture",24    "American craftsman style",25    "American Foursquare architecture",26    "Ancient Egyptian architecture",27    "Art Deco architecture",28    "Art Nouveau architecture",29    "Baroque architecture",30    "Bauhaus architecture",31    "Beaux-Arts architecture",32    "Brutalism architecture",33    "Byzantine architecture",34    "Chicago school architecture",35    "Colonial architecture",36    "Deconstructivism",37    "Edwardian architecture",38    "Georgian architecture",39    "Gothic architecture",40    "Greek Revival architecture",41    "International style",42    "Islamic architecture",43    "Novelty architecture",44    "Palladian architecture",45    "Postmodern architecture",46    "Queen Anne architecture",47    "Romanesque architecture",48    "Russian Revival architecture",49    "Tudor Revival architecture"50]51 52print(len(labels))53model = torchvision.models.efficientnet_v2_l()54 55model.classifier = nn.Sequential(56    nn.Dropout(p=0.4, inplace=True),57    nn.Linear(1280, len(labels), bias=True)58)59 60load_model(model)61 62 63target_layers = model.features[-1]64classifier = model.classifier65cam = LayerCAM(model=model, target_layers=target_layers, use_cuda=False)66dff = DeepFeatureFactorization(67    model=model, target_layer=target_layers, computation_on_concepts=classifier)68 69 70def show_factorization_on_image(img: np.ndarray,71                                explanations: np.ndarray,72                                colors: List[np.ndarray] = None,73                                image_weight: float = 0.5,74                                concept_labels: List = None) -> np.ndarray:75    n_components = explanations.shape[0]76    if colors is None:77        # taken from https://github.com/edocollins/DFF/blob/master/utils.py78        _cmap = plt.cm.get_cmap('gist_rainbow')79        colors = [80            np.array(81                _cmap(i)) for i in np.arange(82                0,83                1,84                1.0 /85                n_components)]86    concept_per_pixel = explanations.argmax(axis=0)87    masks = []88    for i in range(n_components):89        mask = np.zeros(shape=(img.shape[0], img.shape[1], 3))90        mask[:, :, :] = colors[i][:3]91        explanation = explanations[i]92        explanation[concept_per_pixel != i] = 093        mask = np.uint8(mask * 255)94        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2HSV)95        mask[:, :, 2] = np.uint8(255 * explanation)96        mask = cv2.cvtColor(mask, cv2.COLOR_HSV2RGB)97        mask = np.float32(mask) / 25598        masks.append(mask)99 100    mask = np.sum(np.float32(masks), axis=0)101    result = img * image_weight + mask * (1 - image_weight)102    result = np.uint8(result * 255)103 104    if concept_labels is not None:105        px = 1 / plt.rcParams['figure.dpi']  # pixel in inches106        fig = plt.figure(figsize=(result.shape[1] * px, result.shape[0] * px))107        plt.rcParams['legend.fontsize'] = 6 * result.shape[0] / 256108        lw = 5 * result.shape[0] / 256109        lines = [Line2D([0], [0], color=colors[i], lw=lw)110                 for i in range(n_components)]111        plt.legend(lines,112                   concept_labels,113 114                   fancybox=False,115                   shadow=False,116                   frameon=False,117                   loc="center")118 119        plt.tight_layout(pad=0, w_pad=0, h_pad=0)120        plt.axis('off')121        fig.canvas.draw()122        data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)123        plt.close(fig=fig)124        data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))125        data = cv2.resize(data, (result.shape[1], result.shape[0]))126        result = np.vstack((result, data))127    return result128 129 130def create_labels(concept_scores, top_k=2):131    """ Create a list with the image-net category names of the top scoring categories"""132    concept_categories = np.argsort(concept_scores, axis=1)[:, ::-1][:, :top_k]133    concept_labels_topk = []134    for concept_index in range(concept_categories.shape[0]):135        categories = concept_categories[concept_index, :]136        concept_labels = []137        for category in categories:138            score = concept_scores[concept_index, category]139            label = f"{labels[category].split(',')[0]}:{score*100:.2f}%"140            concept_labels.append(label)141        concept_labels_topk.append("\n".join(concept_labels))142    return concept_labels_topk143 144 145def predict(rgb_img, top_k):146    print(top_k)147    inp_01 = transforms.Compose(148        [149            transforms.ToTensor(),150            transforms.Normalize([0.4937, 0.5060, 0.5030], [151                                 0.2705, 0.2653, 0.2998]),152            transforms.Resize((224, 224)),153        ])(rgb_img)154 155    model.eval()156    with torch.no_grad():157        prediction = torch.nn.functional.softmax(158            model(inp_01.unsqueeze(0))[0], dim=0)159        confidences = {labels[i]: float(prediction[i])160                       for i in range(len(labels))}161 162    concepts, batch_explanations, concept_outputs = dff(163        inp_01.unsqueeze(0), 5)164 165    concept_outputs = torch.softmax(166        torch.from_numpy(concept_outputs), axis=-1).numpy()167    concept_label_strings = create_labels(concept_outputs, top_k=top_k)168 169    print(inp_01.shape)170    print(batch_explanations[0].shape)171    res = cv2.resize(np.transpose(172        batch_explanations[0], (1, 2, 0)), (rgb_img.size[0], rgb_img.size[1]))173    res = np.transpose(res, (2, 0, 1))174    print(res.shape)175 176    visualization_01 = show_factorization_on_image(np.float32(rgb_img)/255.0,177                                                   res,178                                                   image_weight=0.3,179                                                   concept_labels=concept_label_strings)180 181    return confidences, visualization_01,182 183 184gr.Interface(fn=predict,185             inputs=[gr.Image(type="pil"), gr.Slider(186                 minimum=1, maximum=4, label="Number of top results", step=1)],187             outputs=[gr.Label(num_top_classes=5), "image"],188             examples=[["./assets/bauhaus.jpg", 1],189                       ["./assets/frank_gehry.jpg", 2], ["./assets/pyramid.jpg", 3]]190             ).launch()191 192 193# examples=["./assets/bauhaus.jpg", "./assets/frank_gehry.jpg", "./assets/pyramid.jpg"]194