def writeimage()

in utils/videowriter.py [0:0]

• 48 lines of code
• 14 McCabe index (conditional complexity)

def writeimage(randid, itemnum, outputs, **kwargs):
    imgout = []

    for k in outputs:
        data = outputs[k]

        if isinstance(data, str):
            img = imgout[-1].permute(1, 2, 0).data.numpy().copy('C')
            cv2.putText(img, data, (20, 60), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 255), 2)
            imgout[-1] = torch.from_numpy(img).permute(2, 0, 1)
        elif len(data.size()) == 3 and data.size(0) == 3: # rgb image
            data = data * torch.tensor(kwargs["colcorrect"][:, None, None])
            imgout.append(data)
        elif len(data.size()) == 3 and data.size(0) == 1: # 1 channel image
            cmap = kwargs["cmap"]
            cmapscale = kwargs["cmapscale"]
            colorbar = kwargs["colorbar"]
            if cmap is None:
                imgout.append(data.repeat(3, 1, 1))
            else:
                cmap = cm.get_cmap(cmap)

                colors = cmap.colors if isinstance(cmap, matplotlib.colors.ListedColormap) else cmap(range(256))
                colors = torch.tensor(colors, device=data.device).permute(1, 0)
                indices = (data[0] * 255. / cmapscale).clamp(min=0., max=255.).long()
                data = colors[:, indices][:3, :, :] * 255.

                if colorbar is not None:
                    data = torch.cat([
                        data,
                        torch.zeros(3, data.size(1), colorbar.size(2), device=data.device)
                        ], dim=2)
                    cbarstarty = data.size(1)//2 - colorbar.size(1)//2
                    data.data[:, cbarstarty:cbarstarty+colorbar.size(1), -colorbar.size(2):] = \
                            colorbar[:3, :, :]

                imgout.append(data)
        elif len(data.size()) == 3 and data.size(0) == 2: # flow field?
            yvals, xvals = torch.meshgrid(
                    torch.linspace(-1., 1., 128, device=data.device),
                    torch.linspace(-1., 1., 128, device=data.device))
            vals = torch.stack([xvals, yvals], axis=-1)

            data[:, :128, :128].data = vals
            imgout.append(flow_compute_color(data))

    imgout = resizecat(imgout, dim=2)

    if imgout.size(2) % 2 != 0:
        imgout = imgout[:, :, :-1]
    if imgout.size(1) % 2 != 0:
        imgout = imgout[:, :-1, :]

    # apply gamma
    imgout = ((imgout / 255.) ** (1. / 1.8) * 255.).clamp(min=0., max=255.).byte()

    # CHW -> HWC
    imgout = imgout.permute(1, 2, 0)

    imgout = imgout.data.to("cpu").numpy()

    Image.fromarray(imgout).save("{}/{:06}.png".format(randid, itemnum))