[FIXED] at::Tensor to UIImage

Issue

I have a PyTorch model and try run it on iOS. I have the next code:

        at::Tensor tensor2 = torch::from_blob(imageBuffer2, {1, 1, 256, 256}, at::kFloat);
        c10::InferenceMode guard;
        auto output = _impl.forward({tensor1, tensor2});
        torch::Tensor tensor_img = output.toTuple()->elements()[0].toTensor();

My question is "How I can convert tensor_img to UIImage?"

I found that functions in PyTorch documentation:

- (UIImage*)convertRGBBufferToUIImage:(unsigned char*)buffer
                            withWidth:(int)width
                           withHeight:(int)height {
    char* rgba = (char*)malloc(width * height * 4);
    for (int i = 0; i < width * height; ++i) {
        rgba[4 * i] = buffer[3 * i];
        rgba[4 * i + 1] = buffer[3 * i + 1];
        rgba[4 * i + 2] = buffer[3 * i + 2];
        rgba[4 * i + 3] = 255;
    }

    size_t bufferLength = width * height * 4;
    CGDataProviderRef provider = CGDataProviderCreateWithData(NULL, rgba, bufferLength, NULL);
    size_t bitsPerComponent = 8;
    size_t bitsPerPixel = 32;
    size_t bytesPerRow = 4 * width;

    CGColorSpaceRef colorSpaceRef = CGColorSpaceCreateDeviceRGB();
    if (colorSpaceRef == NULL) {
        NSLog(@"Error allocating color space");
        CGDataProviderRelease(provider);
        return nil;
    }

    CGBitmapInfo bitmapInfo = kCGBitmapByteOrderDefault | kCGImageAlphaPremultipliedLast;
    CGColorRenderingIntent renderingIntent = kCGRenderingIntentDefault;

    CGImageRef iref = CGImageCreate(width,
        height,
        bitsPerComponent,
        bitsPerPixel,
        bytesPerRow,
        colorSpaceRef,
        bitmapInfo,
        provider,
        NULL,
        YES,
        renderingIntent);

    uint32_t* pixels = (uint32_t*)malloc(bufferLength);

    if (pixels == NULL) {
        NSLog(@"Error: Memory not allocated for bitmap");
        CGDataProviderRelease(provider);
        CGColorSpaceRelease(colorSpaceRef);
        CGImageRelease(iref);
        return nil;
    }

    CGContextRef context = CGBitmapContextCreate(pixels,
        width,
        height,
        bitsPerComponent,
        bytesPerRow,
        colorSpaceRef,
        bitmapInfo);

    if (context == NULL) {
        NSLog(@"Error context not created");
        free(pixels);
    }

    UIImage* image = nil;
    if (context) {
        CGContextDrawImage(context, CGRectMake(0.0f, 0.0f, width, height), iref);
        CGImageRef imageRef = CGBitmapContextCreateImage(context);
        if ([UIImage respondsToSelector:@selector(imageWithCGImage:scale:orientation:)]) {
            float scale = [[UIScreen mainScreen] scale];
            image = [UIImage imageWithCGImage:imageRef scale:scale orientation:UIImageOrientationUp];
        } else {
            image = [UIImage imageWithCGImage:imageRef];
        }

        CGImageRelease(imageRef);
        CGContextRelease(context);
    }

    CGColorSpaceRelease(colorSpaceRef);
    CGImageRelease(iref);
    CGDataProviderRelease(provider);

    if (pixels) {
        free(pixels);
    }
    return image;
}

@end

If I correctly understand, that function can convert unsigned char * to UIImage. I think that I need convert my tensor_img to unsigned char*, but I don't understand how I can do it.

Solution

The 1st code its a torch bridge and 2nd code is UIImage helper which I run from Swift. Anyway, I resolve that issue, we can close it. Code example:

for (int i = 0; i < 3 * width * height; i++) { 
  [results addObject:@(floatBuffer[i])]; 
} 

NSMutableData* data = [NSMutableData dataWithLength:sizeof(float) * 3 * width * height];

float* buffer = (float*)[data mutableBytes]; 

for (int j = 0; j < 3 * width * height; j++) {
  buffer[j] = [results[j] floatValue];
} 
return buffer;

Answered By - MrSandmanRUS