{"cells": [{"cell_type": "markdown", "id": 302934307671667531413257853548643485645, "metadata": {}, "source": ["# Gradio Demo: Echocardiogram-Segmentation"]}, {"cell_type": "code", "execution_count": null, "id": 272996653310673477252411125948039410165, "metadata": {}, "outputs": [], "source": ["!pip install -q gradio -f https://download.pytorch.org/whl/torch_stable.html numpy matplotlib wget torch==1.6.0+cpu  torchvision==0.7.0+cpu   "]}, {"cell_type": "code", "execution_count": null, "id": 288918539441861185822528903084949547379, "metadata": {}, "outputs": [], "source": ["# Downloading files from the demo repo\n", "import os\n", "!wget -q https://github.com/gradio-app/gradio/raw/main/demo/Echocardiogram-Segmentation/img1.jpg\n", "!wget -q https://github.com/gradio-app/gradio/raw/main/demo/Echocardiogram-Segmentation/img2.jpg"]}, {"cell_type": "code", "execution_count": null, "id": 44380577570523278879349135829904343037, "metadata": {}, "outputs": [], "source": ["import os\n", "import numpy as np\n", "import torch\n", "import torchvision\n", "import wget \n", "\n", "\n", "destination_folder = \"output\"\n", "destination_for_weights = \"weights\"\n", "\n", "if os.path.exists(destination_for_weights):\n", "    print(\"The weights are at\", destination_for_weights)\n", "else:\n", "    print(\"Creating folder at \", destination_for_weights, \" to store weights\")\n", "    os.mkdir(destination_for_weights)\n", "    \n", "segmentationWeightsURL = 'https://github.com/douyang/EchoNetDynamic/releases/download/v1.0.0/deeplabv3_resnet50_random.pt'\n", "\n", "if not os.path.exists(os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL))):\n", "    print(\"Downloading Segmentation Weights, \", segmentationWeightsURL,\" to \",os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL)))\n", "    filename = wget.download(segmentationWeightsURL, out = destination_for_weights)\n", "else:\n", "    print(\"Segmentation Weights already present\")\n", "\n", "torch.cuda.empty_cache()\n", "\n", "def collate_fn(x):\n", "    x, f = zip(*x)\n", "    i = list(map(lambda t: t.shape[1], x))\n", "    x = torch.as_tensor(np.swapaxes(np.concatenate(x, 1), 0, 1))\n", "    return x, f, i\n", "\n", "model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=False, aux_loss=False)\n", "model.classifier[-1] = torch.nn.Conv2d(model.classifier[-1].in_channels, 1, kernel_size=model.classifier[-1].kernel_size)\n", "\n", "print(\"loading weights from \", os.path.join(destination_for_weights, \"deeplabv3_resnet50_random\"))\n", "\n", "if torch.cuda.is_available():\n", "    print(\"cuda is available, original weights\")\n", "    device = torch.device(\"cuda\")\n", "    model = torch.nn.DataParallel(model)\n", "    model.to(device)\n", "    checkpoint = torch.load(os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL)))\n", "    model.load_state_dict(checkpoint['state_dict'])\n", "else:\n", "    print(\"cuda is not available, cpu weights\")\n", "    device = torch.device(\"cpu\")\n", "    checkpoint = torch.load(os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL)), map_location = \"cpu\")\n", "    state_dict_cpu = {k[7:]: v for (k, v) in checkpoint['state_dict'].items()}\n", "    model.load_state_dict(state_dict_cpu)\n", "\n", "model.eval()\n", "\n", "def segment(input):\n", "    inp = input\n", "    x = inp.transpose([2, 0, 1])  #  channels-first\n", "    x = np.expand_dims(x, axis=0)  # adding a batch dimension    \n", "    \n", "    mean = x.mean(axis=(0, 2, 3))\n", "    std = x.std(axis=(0, 2, 3))\n", "    x = x - mean.reshape(1, 3, 1, 1)\n", "    x = x / std.reshape(1, 3, 1, 1)\n", "    \n", "    with torch.no_grad():\n", "        x = torch.from_numpy(x).type('torch.FloatTensor').to(device)\n", "        output = model(x)    \n", "    \n", "    y = output['out'].numpy()\n", "    y = y.squeeze()\n", "    \n", "    out = y>0    \n", "    \n", "    mask = inp.copy()\n", "    mask[out] = np.array([0, 0, 255])\n", "    \n", "    return mask\n", "\n", "import gradio as gr\n", "\n", "i = gr.Image(shape=(112, 112), label=\"Echocardiogram\")\n", "o = gr.Image(label=\"Segmentation Mask\")\n", "\n", "examples = [[\"img1.jpg\"], [\"img2.jpg\"]]\n", "title = None #\"Left Ventricle Segmentation\"\n", "description = \"This semantic segmentation model identifies the left ventricle in echocardiogram images.\"\n", "# videos. Accurate evaluation of the motion and size of the left ventricle is crucial for the assessment of cardiac function and ejection fraction. In this interface, the user inputs apical-4-chamber images from echocardiography videos and the model will output a prediction of the localization of the left ventricle in blue. This model was trained on the publicly released EchoNet-Dynamic dataset of 10k echocardiogram videos with 20k expert annotations of the left ventricle and published as part of \u2018Video-based AI for beat-to-beat assessment of cardiac function\u2019 by Ouyang et al. in Nature, 2020.\"\n", "thumbnail = \"https://raw.githubusercontent.com/gradio-app/hub-echonet/master/thumbnail.png\"\n", "gr.Interface(segment, i, o, examples=examples, allow_flagging=False, analytics_enabled=False, thumbnail=thumbnail, cache_examples=False).launch()\n"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}