Initial Space commit

app.py

@@ -0,0 +1,110 @@
+import streamlit as st
+import os
+import urllib
+import fastai.vision.all as fai_vision
+import numpy as np
+from pathlib import Path
+import pathlib
+from PIL import Image
+import platform
+import altair as alt
+import pandas as pd
+
+def main():
+    st.title('Fish Masker and Classifier')
+   
+    data_loader, segmenter = load_unet_model()
+    classification_model = load_classification_model()
+    
+    st.markdown("Upload an Amazonian fish photo for masking.")
+    uploaded_image = st.file_uploader("", IMAGE_TYPES)
+    if uploaded_image:
+        image_data = uploaded_image.read()
+        st.markdown('## Original image')
+        st.image(image_data, use_column_width=True)
+
+        original_pil = Image.open(uploaded_image)
+
+        original_pil.save('original.jpg')
+
+        single_file = [Path('original.jpg')]
+        single_pil = Image.open(single_file[0])
+        input_dl = segmenter.dls.test_dl(single_file)
+        masks, _ = segmenter.get_preds(dl=input_dl)
+        masked_pil, percentage_fish = mask_fish_pil(single_pil, masks[0])
+
+        st.markdown('## Masked image')
+        st.markdown(f'**{percentage_fish:.1f}%** of pixels were labeled as "fish"')
+        st.image(masked_pil, use_column_width=True)
+
+        masked_pil.save('masked.jpg')
+
+        st.markdown('## Classification')
+
+        prediction = classification_model.predict('masked.jpg')
+        pred_chart = predictions_to_chart(prediction, classes = classification_model.dls.vocab)
+        st.altair_chart(pred_chart, use_container_width=True)
+
+
+def mask_fish_pil(unmasked_fish, fastai_mask):
+    unmasked_np = np.array(unmasked_fish)
+    np_mask = fastai_mask.argmax(dim=0).numpy()
+    total_pixels = np_mask.size
+    fish_pixels = np.count_nonzero(np_mask)
+    percentage_fish = (fish_pixels / total_pixels) * 100
+    np_mask = (255 / np_mask.max() * (np_mask - np_mask.min())).astype(np.uint8)
+    np_mask = np.array(Image.fromarray(np_mask).resize(unmasked_np.shape[1::-1], Image.BILINEAR))
+    np_mask = np_mask.reshape(*np_mask.shape, 1) / 255
+    masked_fish_np = (unmasked_np * np_mask).astype(np.uint8)
+    masked_fish_pil = Image.fromarray(masked_fish_np)
+    return masked_fish_pil, percentage_fish
+
+def predictions_to_chart(prediction, classes):
+    pred_rows = []
+    for i, conf in enumerate(list(prediction[2])):
+        pred_row = {'class': classes[i],
+                    'probability': round(float(conf) * 100,2)}
+        pred_rows.append(pred_row)
+    pred_df = pd.DataFrame(pred_rows)
+    pred_df.head()
+    top_probs = pred_df.sort_values('probability', ascending=False).head(4)
+    chart = (
+        alt.Chart(top_probs)
+        .mark_bar()
+        .encode(
+            x=alt.X("probability:Q", scale=alt.Scale(domain=(0, 100))),
+            y=alt.Y("class:N",
+                    sort=alt.EncodingSortField(field="probability", order="descending"))
+        )
+    )
+    return chart
+
+@st.cache(allow_output_mutation=True)
+def load_unet_model():
+    data_loader = fai_vision.SegmentationDataLoaders.from_label_func(
+        path = Path("."),
+        bs = 1,
+        fnames = [Path('test_fish.jpg')],
+        label_func = lambda x: x,
+        codes = np.array(["Photo", "Masks"], dtype=str),
+        item_tfms = [fai_vision.Resize(256, method = 'squish'),],
+        batch_tfms = [fai_vision.IntToFloatTensor(div_mask = 255)],
+        valid_pct = 0.2, num_workers = 0)
+    segmenter = fai_vision.unet_learner(data_loader, fai_vision.resnet34)
+    segmenter.load('fish_mask_model')
+    return data_loader, segmenter
+
+@st.cache(allow_output_mutation=True)
+def load_classification_model():
+    plt = platform.system()
+
+    if plt == 'Linux' or plt == 'Darwin': 
+        pathlib.WindowsPath = pathlib.PosixPath
+    inf_model = fai_vision.load_learner('models/fish_classification_model.pkl', cpu=True)
+
+    return inf_model
+
+IMAGE_TYPES = ["png", "jpg","jpeg"]
+
+if __name__ == "__main__":
+    main()

models/fish_classification_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ac16550590dd60da201ce13e2f1b057d5343ef490db8663c463f8bbefef610e
+size 179319095

models/fish_mask_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29b8afc516eb9f19e99dc53e924839a7157ac241d13f0945aec4717574c7908a
+size 494929527

requirements.txt

@@ -0,0 +1,5 @@
+streamlit==0.89
+fastai==2.2
+protobuf==3.20
+altair
+pandas